Adds support for sending queries using ActiveRecord::Base.connection.execute to slaves

Author: MishaConway

lib/octopus/proxy.rb

@@ -1,9 +1,12 @@
 require 'set'
 require 'octopus/slave_group'
 require 'octopus/load_balancing/round_robin'
+require 'octopus/query_analysis'
 
 module Octopus
   class Proxy
+    include ::Octopus::QueryAnalysis
+
     attr_accessor :config, :sharded
 
     CURRENT_MODEL_KEY = 'octopus.current_model'.freeze
@@ -293,11 +296,11 @@ def method_missing(method, *args, &block)
         self.last_current_shard = current_shard
         clean_connection_proxy
         conn.send(method, *args, &block)
-      elsif should_send_queries_to_shard_slave_group?(method)
+      elsif should_send_queries_to_shard_slave_group?(method, args)
         send_queries_to_shard_slave_group(method, *args, &block)
-      elsif should_send_queries_to_slave_group?(method)
+      elsif should_send_queries_to_slave_group?(method, args)
         send_queries_to_slave_group(method, *args, &block)
-      elsif should_send_queries_to_replicated_databases?(method)
+      elsif should_send_queries_to_replicated_databases?(method, args)
         send_queries_to_selected_slave(method, *args, &block)
       else
         select_connection.send(method, *args, &block)
@@ -337,16 +340,16 @@ def connected?
       @shards.any? { |_k, v| v.connected? }
     end
 
-    def should_send_queries_to_shard_slave_group?(method)
-      should_use_slaves_for_method?(method) && @shards_slave_groups.try(:[], current_shard).try(:[], current_slave_group).present?
+    def should_send_queries_to_shard_slave_group?(method, args)
+      should_use_slaves_for_method?(method, args) && @shards_slave_groups.try(:[], current_shard).try(:[], current_slave_group).present?
     end
 
     def send_queries_to_shard_slave_group(method, *args, &block)
       send_queries_to_balancer(@shards_slave_groups[current_shard][current_slave_group], method, *args, &block)
     end
 
-    def should_send_queries_to_slave_group?(method)
-      should_use_slaves_for_method?(method) && @slave_groups.try(:[], current_slave_group).present?
+    def should_send_queries_to_slave_group?(method, args)
+      should_use_slaves_for_method?(method, args) && @slave_groups.try(:[], current_slave_group).present?
     end
 
     def send_queries_to_slave_group(method, *args, &block)
@@ -431,8 +434,8 @@ def should_clean_connection_proxy?(method)
     end
 
     # Try to use slaves if and only if `replicated: true` is specified in `shards.yml` and no slaves groups are defined
-    def should_send_queries_to_replicated_databases?(method)
-      @replicated && method.to_s =~ /select/ && !block && !slaves_grouped?
+    def should_send_queries_to_replicated_databases?(method, args)
+      @replicated && select?(method, args) && !block && !slaves_grouped?
     end
 
     def current_model_replicated?
@@ -458,8 +461,20 @@ def send_queries_to_selected_slave(method, *args, &block)
     # (3) It's a SELECT query
     # while ensuring that we revert `current_shard` from the selected slave to the (shard's) master
     # not to make queries other than SELECT leak to the slave.
-    def should_use_slaves_for_method?(method)
-      current_model_replicated? && method.to_s =~ /select/
+    def should_use_slaves_for_method?(method, args)
+      current_model_replicated? && select?(method, args)
+    end
+
+    # Given an ActiveRecord::Base.connection method and its arguments, determine if it is a select query.
+    def select?(method, args)
+      is_select = method.to_s =~ /select/
+      unless is_select
+        query = args.first
+        if query.kind_of? String
+          is_select = method.to_s =~ /execute/ && definitely_select_query?(query) && !possibly_multiple_queries?(query)
+        end
+      end
+      is_select
     end
 
     def slaves_grouped?

lib/octopus/query_analysis.rb

@@ -0,0 +1,16 @@
+module Octopus
+  module QueryAnalysis
+    # Given a mysql query string, determines if it is definitely a select query. Due to the simple regex used, it will
+    # sometimes miss detecting valid select queries, hence why it only determines if something is definitely a select.
+    def definitely_select_query?( str )
+      str =~ /^\s*select/i
+    end
+
+    # Given a mysql query string, determines if the string might contain multiple queries.
+    # We are simply checking if it contains a semi colon with non whitespace to the right of it, so this check will
+    # sometimes falsely detect a string containing one query as sometimes having multiple queries.
+    def possibly_multiple_queries?( str )
+      str =~ /;.*\S+.*$/
+    end
+  end
+end