diff options
Diffstat (limited to 'activerecord/lib/active_record/identity_map.rb')
-rw-r--r-- | activerecord/lib/active_record/identity_map.rb | 24 |
1 files changed, 21 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 |