diff options
author | Aaron Patterson <aaron.patterson@gmail.com> | 2011-05-03 09:52:41 -0700 |
---|---|---|
committer | Aaron Patterson <aaron.patterson@gmail.com> | 2011-05-03 09:59:22 -0700 |
commit | d0719c59384c387aebfdc4a78f6e985bcf5bb270 (patch) | |
tree | 84430a5ee109fbb3bdbd7a2ec2f7753cb786d116 | |
parent | 3f897c1a4cef0cb9b53234e10d10d4b8744d61b0 (diff) | |
download | rails-d0719c59384c387aebfdc4a78f6e985bcf5bb270.tar.gz rails-d0719c59384c387aebfdc4a78f6e985bcf5bb270.tar.bz2 rails-d0719c59384c387aebfdc4a78f6e985bcf5bb270.zip |
proxying the body in the IM middleware so that IM is available for streaming ERb
-rw-r--r-- | activerecord/lib/active_record/identity_map.rb | 24 | ||||
-rw-r--r-- | activerecord/test/cases/identity_map/middleware_test.rb | 26 |
2 files changed, 47 insertions, 3 deletions
diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb index 4ad0b0d205..647f9e11d2 100644 --- a/activerecord/lib/active_record/identity_map.rb +++ b/activerecord/lib/active_record/identity_map.rb @@ -98,14 +98,32 @@ module ActiveRecord end class Middleware + class Body #:nodoc: + def initialize(target, original) + @target = target + @original = original + end + + def each(&block) + @target.each(&block) + end + + def close + @target.close if @target.respond_to?(:close) + ensure + IdentityMap.enabled = @original + end + end + def initialize(app) @app = app end def call(env) - ActiveRecord::IdentityMap.use do - @app.call(env) - end + enabled = IdentityMap.enabled + IdentityMap.enabled = true + status, headers, body = @app.call(env) + [status, headers, Body.new(body, enabled)] end end end diff --git a/activerecord/test/cases/identity_map/middleware_test.rb b/activerecord/test/cases/identity_map/middleware_test.rb index 1004e958e5..130f370da6 100644 --- a/activerecord/test/cases/identity_map/middleware_test.rb +++ b/activerecord/test/cases/identity_map/middleware_test.rb @@ -29,6 +29,32 @@ module ActiveRecord } mw.call({}) end + + class Enum < Struct.new(:iter) + def each(&b) + iter.call(&b) + end + end + + def test_im_enabled_during_body_each + mw = Middleware.new lambda { |env| + [200, {}, Enum.new(lambda { |&b| + assert IdentityMap.enabled?, 'identity map should be enabled' + b.call "hello" + })] + } + body = mw.call({}).last + body.each { |x| assert_equal 'hello', x } + end + + def test_im_disabled_after_body_close + mw = Middleware.new lambda { |env| [200, {}, []] } + body = mw.call({}).last + assert IdentityMap.enabled?, 'identity map should be enabled' + body.close + assert !IdentityMap.enabled?, 'identity map should be disabled' + end + end end end |