From d222f7de572f28fc2fa185e9f21cac6f7e6c84f0 Mon Sep 17 00:00:00 2001
From: Pratik Naik <pratiknaik@gmail.com>
Date: Mon, 12 Oct 2015 12:33:48 -0500
Subject: Make sure active record queries are logged with the current
 connection tags

---
 lib/action_cable/server/worker.rb                  |  7 ++++++-
 .../worker/active_record_connection_management.rb  | 22 ++++++++++++++++++++++
 .../server/worker/clear_database_connections.rb    | 22 ----------------------
 3 files changed, 28 insertions(+), 23 deletions(-)
 create mode 100644 lib/action_cable/server/worker/active_record_connection_management.rb
 delete mode 100644 lib/action_cable/server/worker/clear_database_connections.rb

(limited to 'lib/action_cable/server')

diff --git a/lib/action_cable/server/worker.rb b/lib/action_cable/server/worker.rb
index d7823ecf93..91496775b8 100644
--- a/lib/action_cable/server/worker.rb
+++ b/lib/action_cable/server/worker.rb
@@ -5,10 +5,13 @@ module ActionCable
       include ActiveSupport::Callbacks
       include Celluloid
 
+      attr_reader :connection
       define_callbacks :work
-      include ClearDatabaseConnections
+      include ActiveRecordConnectionManagement
 
       def invoke(receiver, method, *args)
+        @connection = receiver
+
         run_callbacks :work do
           receiver.send method, *args
         end
@@ -20,6 +23,8 @@ module ActionCable
       end
 
       def run_periodic_timer(channel, callback)
+        @connection = channel.connection
+
         run_callbacks :work do
           callback.respond_to?(:call) ? channel.instance_exec(&callback) : channel.send(callback)
         end
diff --git a/lib/action_cable/server/worker/active_record_connection_management.rb b/lib/action_cable/server/worker/active_record_connection_management.rb
new file mode 100644
index 0000000000..1ede0095f8
--- /dev/null
+++ b/lib/action_cable/server/worker/active_record_connection_management.rb
@@ -0,0 +1,22 @@
+module ActionCable
+  module Server
+    class Worker
+      # Clear active connections between units of work so the long-running channel or connection processes do not hoard connections.
+      module ActiveRecordConnectionManagement
+        extend ActiveSupport::Concern
+
+        included do
+          if defined?(ActiveRecord::Base)
+            set_callback :work, :around, :with_database_connections
+          end
+        end
+
+        def with_database_connections
+          ActiveRecord::Base.logger.tagged(*connection.logger.tags) { yield }
+        ensure
+          ActiveRecord::Base.clear_active_connections!
+        end
+      end
+    end
+  end
+end
\ No newline at end of file
diff --git a/lib/action_cable/server/worker/clear_database_connections.rb b/lib/action_cable/server/worker/clear_database_connections.rb
deleted file mode 100644
index 722d363a41..0000000000
--- a/lib/action_cable/server/worker/clear_database_connections.rb
+++ /dev/null
@@ -1,22 +0,0 @@
-module ActionCable
-  module Server
-    class Worker
-      # Clear active connections between units of work so the long-running channel or connection processes do not hoard connections.
-      module ClearDatabaseConnections
-        extend ActiveSupport::Concern
-
-        included do
-          if defined?(ActiveRecord::Base)
-            set_callback :work, :around, :with_database_connections
-          end
-        end
-
-        def with_database_connections
-          yield
-        ensure
-          ActiveRecord::Base.clear_active_connections!
-        end
-      end
-    end
-  end
-end
\ No newline at end of file
-- 
cgit v1.2.3