From faa8a87a3fd94018f0648c8bafffa979ce8261e4 Mon Sep 17 00:00:00 2001
From: Maxim Polunin <maks.polunin@gmail.com>
Date: Sat, 20 Feb 2016 13:09:27 +0300
Subject: [PATCH] Use scoped subscription to ActiveSupport::Notifications

---
 lib/chewy/rake_helper.rb |  7 +++++--
 lib/tasks/chewy.rake     | 32 ++++++++++++++++++--------------
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/lib/chewy/rake_helper.rb b/lib/chewy/rake_helper.rb
index 452bb2a09..3dcbbe7ee 100644
--- a/lib/chewy/rake_helper.rb
+++ b/lib/chewy/rake_helper.rb
@@ -2,8 +2,8 @@ module Chewy
   module RakeHelper
     class << self
 
-      def subscribe_task_stats!
-        ActiveSupport::Notifications.subscribe('import_objects.chewy') do |name, start, finish, id, payload|
+      def subscribed_task_stats(&block)
+        callback = ->(name, start, finish, id, payload) do
           duration = (finish - start).round(2)
           puts "  Imported #{payload[:type]} for #{duration}s, documents total: #{payload[:import].try(:[], :index).to_i}"
           payload[:errors].each do |action, errors|
@@ -14,6 +14,9 @@ def subscribe_task_stats!
             end
           end if payload[:errors]
         end
+        ActiveSupport::Notifications.subscribed(callback, 'import_objects.chewy') do
+          yield
+        end
       end
 
       def eager_load_chewy!
diff --git a/lib/tasks/chewy.rake b/lib/tasks/chewy.rake
index f014962cb..261ef1c49 100644
--- a/lib/tasks/chewy.rake
+++ b/lib/tasks/chewy.rake
@@ -3,39 +3,43 @@ require 'chewy/rake_helper'
 namespace :chewy do
   desc 'Destroy, recreate and import data to specified index'
   task :reset, [:index] => :environment do |task, args|
-    Chewy::RakeHelper.subscribe_task_stats!
+    Chewy::RakeHelper.subscribed_task_stats do
 
-    if args[:index].present?
-      Chewy::RakeHelper.reset_index(args[:index])
-    else
-      Chewy::RakeHelper.reset_all
+      if args[:index].present?
+        Chewy::RakeHelper.reset_index(args[:index])
+      else
+        Chewy::RakeHelper.reset_all
+      end
     end
   end
 
   namespace :reset do
     desc 'Destroy, recreate and import data for all found indexes'
     task all: :environment do
-      Chewy::RakeHelper.subscribe_task_stats!
-      Chewy::RakeHelper.reset_all
+      Chewy::RakeHelper.subscribed_task_stats do
+        Chewy::RakeHelper.reset_all
+      end
     end
   end
 
   desc 'Updates data specified index'
   task :update, [:index] => :environment do |task, args|
-    Chewy::RakeHelper.subscribe_task_stats!
+    Chewy::RakeHelper.subscribed_task_stats do
 
-    if args[:index].present?
-      Chewy::RakeHelper.update_index(args[:index])
-    else
-      Chewy::RakeHelper.update_all
+      if args[:index].present?
+        Chewy::RakeHelper.update_index(args[:index])
+      else
+        Chewy::RakeHelper.update_all
+      end
     end
   end
 
   namespace :update do
     desc 'Updates data for all found indexes'
     task all: :environment do
-      Chewy::RakeHelper.subscribe_task_stats!
-      Chewy::RakeHelper.update_all
+      Chewy::RakeHelper.subscribed_task_stats do
+        Chewy::RakeHelper.update_all
+      end
     end
   end
 end