Initial implementation of a "cache accumulator" behavior.

Mix Puppet::Util::CacheAccumulator into an ActiveRecord-like class, and then for any
attribute in that class on which you are likely to call find_or_create_by_*, specify:

 accumulates :foo

and instead of :find_or_create_by_foo use :accumulate_by_foo.

The class will cache known results keyed by values of :foo.

Do an initial bulk-lookup: class.accumulate_by_foo('foo1', 'foo2', 'foo3', 'foo4')

2 files changed, 134 insertions, 0 deletions

diff --git a/lib/puppet/util/rails/cache_accumulator.rb b/lib/puppet/util/rails/cache_accumulator.rb
new file mode 100644
index 000000000..7c59a5d3f
--- /dev/null
+++ b/lib/puppet/util/rails/cache_accumulator.rb
@@ -0,0 +1,65 @@
+require 'puppet/util'
+
+module Puppet::Util::CacheAccumulator
+    def self.included(klass)
+        klass.extend ClassMethods
+    end
+
+    class Base
+        attr_reader :klass, :attribute
+
+        def initialize(klass, attribute)
+            @klass = klass
+            @attribute = attribute
+            @find_or_create = "find_or_create_by_#{@attribute.to_s}".intern
+        end
+
+        def store
+            @store || reset
+        end
+
+        def reset
+            @store = {}
+        end
+
+        def find(*keys)
+            result = nil
+            if keys.length == 1
+                result = store[keys[0]] ||= @klass.send(@find_or_create, *keys)
+            else
+                found, missing = keys.partition {|k| store.include? k}
+                result = found.length
+                result += do_multi_find(missing) if missing.length > 0
+            end
+            result
+        end
+
+        def do_multi_find(keys)
+            result = 0
+            @klass.find(:all, :conditions => {@attribute => keys}).each do |obj|
+                store[obj.send(@attribute)] = obj
+                result += 1
+            end
+            result
+        end
+    end
+
+    module ClassMethods
+        def accumulates(*attributes)
+            attributes.each {|attrib| install_accumulator(attrib)}
+        end
+
+        def accumulators
+            @accumulators ||= {}
+        end
+
+        def install_accumulator(attribute)
+            self.accumulators[attribute] = Base.new(self, attribute)
+            module_eval %{
+                def self.accumulate_by_#{attribute.to_s}(*keys)
+                    accumulators[:#{attribute.to_s}].find(*keys)
+                end
+            }
+        end
+    end
+end
diff --git a/spec/unit/util/cache_accumulator.rb b/spec/unit/util/cache_accumulator.rb
new file mode 100644
index 000000000..55ecb7e34
--- /dev/null
+++ b/spec/unit/util/cache_accumulator.rb
@@ -0,0 +1,69 @@
+#!/usr/bin/env ruby
+
+require File.dirname(__FILE__) + '/../../spec_helper'
+require 'puppet/util/rails/cache_accumulator'
+
+describe Puppet::Util::CacheAccumulator do
+    before :each do
+        @test_class = Class.new do
+            attr_accessor :name
+
+            include Puppet::Util::CacheAccumulator
+            accumulates :name
+
+            def initialize(n)
+                self.name = n
+            end
+        end
+    end
+
+    it 'should delegate to underlying find_or_create_by_* method and accumulate results' do
+        @test_class.expects(:find_or_create_by_name).with('foo').returns(@test_class.new('foo')).once
+        obj = @test_class.accumulate_by_name('foo')
+        obj.name.should == 'foo'
+        @test_class.accumulate_by_name('foo').should == obj
+    end
+
+    it 'should delegate bulk lookups to find with appropriate arguments and returning result count' do
+        @test_class.expects(:find).with(:all,
+                                        :conditions => {:name => ['a', 'b', 'c']}
+                                       ).returns(['a','b','c'].collect {|n| @test_class.new(n)}).once
+        @test_class.accumulate_by_name('a', 'b', 'c').should == 3
+    end
+
+    it 'should only need find_or_create_by_name lookup for missing bulk entries' do
+        @test_class.expects(:find).with(:all,
+                                        :conditions => {:name => ['a', 'b']}
+                                       ).returns([ @test_class.new('a') ]).once
+        @test_class.expects(:find_or_create_by_name).with('b').returns(@test_class.new('b')).once
+        @test_class.expects(:find_or_create_by_name).with('a').never
+        @test_class.accumulate_by_name('a','b').should == 1
+        @test_class.accumulate_by_name('a').name.should == 'a'
+        @test_class.accumulate_by_name('b').name.should == 'b'
+    end
+
+    it 'should keep consumer classes separate' do
+        @alt_class = Class.new do
+            attr_accessor :name
+
+            include Puppet::Util::CacheAccumulator
+            accumulates :name
+
+            def initialize(n)
+                self.name = n
+            end
+        end
+        name = 'foo'
+        @test_class.expects(:find_or_create_by_name).with(name).returns(@test_class.new(name)).once
+        @alt_class.expects(:find_or_create_by_name).with(name).returns(@alt_class.new(name)).once
+        
+        [@test_class, @alt_class].each do |klass|
+            klass.accumulate_by_name(name).name.should == name
+            klass.accumulate_by_name(name).class.should == klass
+        end
+    end
+
+    it 'should clear accumulated cache with reset_*_accumulator' do
+        # figure out how to test this appropriately...
+    end
+end