Added integration with ActiveRecord instrumentation for better loggin…

…g support for Rails 3+

lib/octopus.rb

@@ -90,6 +90,8 @@ def self.using(shard, &block)
   require "octopus/rails3/association"
   require "octopus/rails3/persistence"
   require "octopus/rails3/arel"
+  require "octopus/rails3/log_subscriber"
+  require "octopus/rails3/abstract_adapter"
 else
   require "octopus/rails2/association"
   require "octopus/rails2/persistence"

lib/octopus/proxy.rb

@@ -1,7 +1,7 @@
 require "set"
 
 class Octopus::Proxy
-  attr_accessor :current_model, :current_shard, :current_group, :block, 
+  attr_accessor :current_model, :current_shard, :current_group, :block,
       :using_enabled, :last_current_shard, :config
 
   def initialize(config)
@@ -40,7 +40,8 @@ def initialize_shards(config)
         value.each do |k, v|
           raise "You have duplicated shard names!" if @shards.has_key?(k.to_sym)
           initialize_adapter(v['adapter'])
-          @shards[k.to_sym] = connection_pool_for(v, "#{v['adapter']}_connection")
+          config_with_octopus_shard = v.merge(:octopus_shard => k)
+          @shards[k.to_sym] = connection_pool_for(config_with_octopus_shard, "#{v['adapter']}_connection")
           @groups[key.to_sym] << k.to_sym
         end
       end
@@ -89,7 +90,7 @@ def select_connection
     # connection pool.  Octopus can potentially retain a reference to a closed
     # connection pool.  Previously, that would work since the pool would just
     # reconnect, but in Rails 3.1 the flag prevents this.
-    if Octopus.rails31? 
+    if Octopus.rails31?
       if !@shards[shard_name].automatic_reconnect
         @shards[shard_name].automatic_reconnect = true
       end
@@ -108,7 +109,7 @@ def should_clean_table_name?
   def run_queries_on_shard(shard, &block)
     older_shard = self.current_shard
     last_block = self.block
-        
+
     begin
       self.block = true
       self.current_shard = shard
@@ -164,7 +165,7 @@ def method_missing(method, *args, &block)
   def respond_to?(method, include_private = false)
     super || select_connection.respond_to?(method, include_private)
   end
-  
+
   def connection_pool
     return @shards[current_shard]
   end
@@ -173,7 +174,7 @@ def connection_pool
   def connection_pool_for(adapter, config)
     ActiveRecord::ConnectionAdapters::ConnectionPool.new(ActiveRecord::Base::ConnectionSpecification.new(adapter, config))
   end
-  
+
   def initialize_adapter(adapter)
     @adapters << adapter
     begin

lib/octopus/rails3/abstract_adapter.rb

@@ -0,0 +1,39 @@
+# Implementation courtesy of db-charmer.
+module Octopus
+  module AbstractAdapter
+    module OctopusShard
+
+      class InstrumenterDecorator < ActiveSupport::BasicObject
+        def initialize(adapter, instrumenter)
+          @adapter = adapter
+          @instrumenter = instrumenter
+        end
+
+        def instrument(name, payload = {}, &block)
+          payload[:octopus_shard] ||= @adapter.octopus_shard
+          @instrumenter.instrument(name, payload, &block)
+        end
+
+        def method_missing(meth, *args, &block)
+          @instrumenter.send(meth, *args, &block)
+        end
+      end
+
+      def self.included(base)
+        base.alias_method_chain :initialize, :octopus_shard
+      end
+
+      def octopus_shard
+        @config[:octopus_shard]
+      end
+
+      def initialize_with_octopus_shard(*args)
+        initialize_without_octopus_shard(*args)
+        @instrumenter = InstrumenterDecorator.new(self, @instrumenter)
+      end
+
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, Octopus::AbstractAdapter::OctopusShard)

lib/octopus/rails3/log_subscriber.rb

@@ -0,0 +1,22 @@
+# Implementation courtesy of db-charmer.
+module Octopus
+  module LogSubscriber
+    def self.included(base)
+      base.send(:attr_accessor, :octopus_shard)
+      base.alias_method_chain :sql, :octopus_shard
+      base.alias_method_chain :debug, :octopus_shard
+    end
+
+    def sql_with_octopus_shard(event)
+      self.octopus_shard = event.payload[:octopus_shard]
+      sql_without_octopus_shard(event)
+    end
+
+    def debug_with_octopus_shard(msg)
+      conn = octopus_shard ? color("[Shard: #{octopus_shard}]", ActiveSupport::LogSubscriber::GREEN, true) : ''
+      debug_without_octopus_shard(conn + msg)
+    end
+  end
+end
+
+ActiveRecord::LogSubscriber.send(:include, Octopus::LogSubscriber)

spec/octopus/log_subscriber_spec.rb

@@ -0,0 +1,21 @@
+require File.expand_path(File.dirname(__FILE__) + '/../spec_helper')
+
+if Octopus.rails3?
+  describe Octopus::LogSubscriber do
+    before :each do
+      @out = StringIO.new
+      @log = Logger.new(@out)
+      ActiveRecord::Base.logger = @log
+      ActiveRecord::Base.logger.level = Logger::DEBUG
+    end
+
+    after :each do
+      ActiveRecord::Base.logger = nil
+    end
+
+    it "should add to the default logger what shard the query was sent" do
+      User.using(:canada).create!(:name => "test")
+      @out.string.should =~ /Shard: canada/
+    end
+  end
+end