ar-octopus 0.8.5 → 0.10.2

data/lib/octopus/proxy.rb

@@ -4,196 +4,64 @@ require 'octopus/load_balancing/round_robin'
 
 module Octopus
   class Proxy
-    attr_accessor :config, :sharded
+    attr_accessor :proxy_config
+
+    delegate :current_model, :current_model=,
+             :current_shard, :current_shard=,
+             :current_group, :current_group=,
+             :current_slave_group, :current_slave_group=,
+             :current_load_balance_options, :current_load_balance_options=,
+             :block, :block=, :fully_replicated?, :has_group?,
+             :shard_names, :shards_for_group, :shards, :sharded, :slaves_list,
+             :shards_slave_groups, :slave_groups, :replicated, :slaves_load_balancer,
+             :config, :initialize_shards, :shard_name, to: :proxy_config, prefix: false
 
     def initialize(config = Octopus.config)
-      initialize_shards(config)
-      initialize_replication(config) if !config.nil? && config['replicated']
+      self.proxy_config = Octopus::ProxyConfig.new(config)
     end
 
-    def initialize_shards(config)
-      @shards = HashWithIndifferentAccess.new
-      @shards_slave_groups = HashWithIndifferentAccess.new
-      @slave_groups = HashWithIndifferentAccess.new
-      @groups = {}
-      @adapters = Set.new
-      @config = ActiveRecord::Base.connection_pool_without_octopus.spec.config
+    # Rails Connection Methods - Those methods are overriden to add custom behavior that helps
+    # Octopus introduce Sharding / Replication.
+    delegate :adapter_name, :add_transaction_record, :case_sensitive_modifier,
+      :type_cast, :to_sql, :quote, :quote_column_name, :quote_table_name,
+      :quote_table_name_for_assignment, :supports_migrations?, :table_alias_for,
+      :table_exists?, :in_clause_length, :supports_ddl_transactions?,
+      :sanitize_limit, :prefetch_primary_key?, :current_database,
+      :combine_bind_parameters, :empty_insert_statement_value, :assume_migrated_upto_version,
+      :schema_cache, :substitute_at, :internal_string_options_for_primary_key, :lookup_cast_type_from_column,
+      :supports_advisory_locks?, :get_advisory_lock, :initialize_internal_metadata_table,
+      :release_advisory_lock, :prepare_binds_for_database, :cacheable_query, :column_name_for_operation,
+      :prepared_statements, :transaction_state, :create_table, to: :select_connection
 
-      unless config.nil?
-        @entire_sharded = config['entire_sharded']
-        @shards_config = config[Octopus.rails_env]
-      end
-
-      @shards_config ||= []
-
-      @shards_config.each do |key, value|
-        if value.is_a?(String)
-          value = resolve_string_connection(value).merge(:octopus_shard => key)
-          initialize_adapter(value['adapter'])
-          @shards[key.to_sym] = connection_pool_for(value, "#{value['adapter']}_connection")
-        elsif value.is_a?(Hash) && value.key?('adapter')
-          value.merge!(:octopus_shard => key)
-          initialize_adapter(value['adapter'])
-          @shards[key.to_sym] = connection_pool_for(value, "#{value['adapter']}_connection")
-
-          slave_group_configs = value.select do |_k, v|
-            structurally_slave_group? v
-          end
-
-          if slave_group_configs.present?
-            slave_groups = HashWithIndifferentAccess.new
-            slave_group_configs.each do |slave_group_name, slave_configs|
-              slaves = HashWithIndifferentAccess.new
-              slave_configs.each do |slave_name, slave_config|
-                @shards[slave_name.to_sym] = connection_pool_for(slave_config, "#{value['adapter']}_connection")
-                slaves[slave_name.to_sym] = slave_name.to_sym
-              end
-              slave_groups[slave_group_name.to_sym] = Octopus::SlaveGroup.new(slaves)
-            end
-            @shards_slave_groups[key.to_sym] = slave_groups
-            @sharded = true
-          end
-        elsif value.is_a?(Hash)
-          @groups[key.to_s] = []
-
-          value.each do |k, v|
-            fail 'You have duplicated shard names!' if @shards.key?(k.to_sym)
-
-            initialize_adapter(v['adapter'])
-            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_s] << k.to_sym
-          end
-
-          if structurally_slave_group? value
-            slaves = Hash[@groups[key.to_s].map { |v| [v, v] }]
-            @slave_groups[key.to_sym] = Octopus::SlaveGroup.new(slaves)
-          end
-        end
-      end
-
-      @shards[:master] ||= ActiveRecord::Base.connection_pool_without_octopus
-    end
-
-    def initialize_replication(config)
-      @replicated = true
-      if config.key?('fully_replicated')
-        @fully_replicated = config['fully_replicated']
-      else
-        @fully_replicated = true
-      end
-
-      @slaves_list = @shards.keys.map(&:to_s).sort
-      @slaves_list.delete('master')
-      @slaves_load_balancer = Octopus::LoadBalancing::RoundRobin.new(@slaves_list)
-    end
-
-    def current_model
-      Thread.current['octopus.current_model']
-    end
-
-    def current_model=(model)
-      Thread.current['octopus.current_model'] = model.is_a?(ActiveRecord::Base) ? model.class : model
-    end
-
-    def current_shard
-      Thread.current['octopus.current_shard'] ||= :master
-    end
-
-    def current_shard=(shard_symbol)
-      self.current_slave_group = nil
-      if shard_symbol.is_a?(Array)
-        shard_symbol.each { |symbol| fail "Nonexistent Shard Name: #{symbol}" if @shards[symbol].nil? }
-      elsif shard_symbol.is_a?(Hash)
-        hash = shard_symbol
-        shard_symbol = hash[:shard]
-        slave_group_symbol = hash[:slave_group]
-
-        if shard_symbol.nil? && slave_group_symbol.nil?
-          fail 'Neither shard or slave group must be specified'
-        end
-
-        if shard_symbol.present?
-          fail "Nonexistent Shard Name: #{shard_symbol}" if @shards[shard_symbol].nil?
-        end
-
-        if slave_group_symbol.present?
-          if (@shards_slave_groups.try(:[], shard_symbol).present? && @shards_slave_groups[shard_symbol][slave_group_symbol].nil?) ||
-              (@shards_slave_groups.try(:[], shard_symbol).nil? && @slave_groups[slave_group_symbol].nil?)
-            fail "Nonexistent Slave Group Name: #{slave_group_symbol} in shards config: #{@shards_config.inspect}"
-          end
-          self.current_slave_group = slave_group_symbol
-        end
-      else
-        fail "Nonexistent Shard Name: #{shard_symbol}" if @shards[shard_symbol].nil?
-      end
-
-      Thread.current['octopus.current_shard'] = shard_symbol
+    def execute(sql, name = nil)
+      conn = select_connection
+      clean_connection_proxy if should_clean_connection_proxy?('execute')
+      conn.execute(sql, name)
     end
 
-    def current_group
-      Thread.current['octopus.current_group']
+    def insert(arel, name = nil, pk = nil, id_value = nil, sequence_name = nil, binds = [])
+      conn = select_connection
+      clean_connection_proxy if should_clean_connection_proxy?('insert')
+      conn.insert(arel, name, pk, id_value, sequence_name, binds)
     end
 
-    def current_group=(group_symbol)
-      # TODO: Error message should include all groups if given more than one bad name.
-      [group_symbol].flatten.compact.each do |group|
-        fail "Nonexistent Group Name: #{group}" unless has_group?(group)
-      end
-
-      Thread.current['octopus.current_group'] = group_symbol
+    def update(arel, name = nil, binds = [])
+      conn = select_connection
+      # Call the legacy should_clean_connection_proxy? method here, emulating an insert.
+      clean_connection_proxy if should_clean_connection_proxy?('insert')
+      conn.update(arel, name, binds)
     end
 
-    def current_slave_group
-      Thread.current['octopus.current_slave_group']
+    def delete(*args, &block)
+      legacy_method_missing_logic('delete', *args, &block)
     end
 
-    def current_slave_group=(slave_group_symbol)
-      Thread.current['octopus.current_slave_group'] = slave_group_symbol
+    def select_all(*args, &block)
+      legacy_method_missing_logic('select_all', *args, &block)
     end
 
-    def block
-      Thread.current['octopus.block']
-    end
-
-    def block=(block)
-      Thread.current['octopus.block'] = block
-    end
-
-    def last_current_shard
-      Thread.current['octopus.last_current_shard']
-    end
-
-    def last_current_shard=(last_current_shard)
-      Thread.current['octopus.last_current_shard'] = last_current_shard
-    end
-
-    def fully_replicated?
-      @fully_replicated || Thread.current['octopus.fully_replicated']
-    end
-
-    # Public: Whether or not a group exists with the given name converted to a
-    # string.
-    #
-    # Returns a boolean.
-    def has_group?(group)
-      @groups.key?(group.to_s)
-    end
-
-    # Public: Retrieves names of all loaded shards.
-    #
-    # Returns an array of shard names as symbols
-    def shard_names
-      @shards.keys
-    end
-
-    # Public: Retrieves the defined shards for a given group.
-    #
-    # Returns an array of shard names as symbols or nil if the group is not
-    # defined.
-    def shards_for_group(group)
-      @groups.fetch(group.to_s, nil)
+    def select_value(*args, &block)
+      legacy_method_missing_logic('select_value', *args, &block)
     end
 
     # Rails 3.1 sets automatic_reconnect to false when it removes
@@ -202,22 +70,14 @@ module Octopus
     # reconnect, but in Rails 3.1 the flag prevents this.
     def safe_connection(connection_pool)
       connection_pool.automatic_reconnect ||= true
-      if !connection_pool.connected? && @shards[:master].connection.query_cache_enabled
+      if !connection_pool.connected? && shards[Octopus.master_shard].connection.query_cache_enabled
         connection_pool.connection.enable_query_cache!
       end
       connection_pool.connection
     end
 
     def select_connection
-      safe_connection(@shards[shard_name])
-    end
-
-    def shard_name
-      current_shard.is_a?(Array) ? current_shard.first : current_shard
-    end
-
-    def should_clean_table_name?
-      @adapters.size > 1
+      safe_connection(shards[shard_name])
     end
 
     def run_queries_on_shard(shard, &_block)
@@ -229,13 +89,23 @@ module Octopus
     end
 
     def send_queries_to_multiple_shards(shards, &block)
-      shards.each do |shard|
+      shards.map do |shard|
         run_queries_on_shard(shard, &block)
       end
     end
 
+    def send_queries_to_group(group, &block)
+      using_group(group) do
+        send_queries_to_multiple_shards(shards_for_group(group), &block)
+      end
+    end
+
+    def send_queries_to_all_shards(&block)
+      send_queries_to_multiple_shards(shard_names.uniq { |shard_name| shards[shard_name] }, &block)
+    end
+
     def clean_connection_proxy
-      self.current_shard = :master
+      self.current_shard = Octopus.master_shard
       self.current_model = nil
       self.current_group = nil
       self.block = nil
@@ -249,7 +119,7 @@ module Octopus
 
     def transaction(options = {}, &block)
       if !sharded && current_model_replicated?
-        run_queries_on_shard(:master) do
+        run_queries_on_shard(Octopus.master_shard) do
           select_connection.transaction(options, &block)
         end
       else
@@ -258,20 +128,7 @@ module Octopus
     end
 
     def method_missing(method, *args, &block)
-      if should_clean_connection_proxy?(method)
-        conn = select_connection
-        self.last_current_shard = current_shard
-        clean_connection_proxy
-        conn.send(method, *args, &block)
-      elsif should_send_queries_to_shard_slave_group?(method)
-        send_queries_to_shard_slave_group(method, *args, &block)
-      elsif should_send_queries_to_slave_group?(method)
-        send_queries_to_slave_group(method, *args, &block)
-      elsif should_send_queries_to_replicated_databases?(method)
-        send_queries_to_selected_slave(method, *args, &block)
-      else
-        select_connection.send(method, *args, &block)
-      end
+      legacy_method_missing_logic(method, *args, &block)
     end
 
     def respond_to?(method, include_private = false)
@@ -279,16 +136,18 @@ module Octopus
     end
 
     def connection_pool
-      @shards[current_shard]
+      shards[current_shard]
     end
 
-    def enable_query_cache!
-      clear_query_cache
-      with_each_healthy_shard { |v| v.connected? && safe_connection(v).enable_query_cache! }
-    end
+    if Octopus.rails4?
+      def enable_query_cache!
+        clear_query_cache
+        with_each_healthy_shard { |v| v.connected? && safe_connection(v).enable_query_cache! }
+      end
 
-    def disable_query_cache!
-      with_each_healthy_shard { |v| v.connected? && safe_connection(v).disable_query_cache! }
+      def disable_query_cache!
+        with_each_healthy_shard { |v| v.connected? && safe_connection(v).disable_query_cache! }
+      end
     end
 
     def clear_query_cache
@@ -301,33 +160,81 @@ module Octopus
 
     def clear_all_connections!
       with_each_healthy_shard(&:disconnect!)
+
+      if Octopus.atleast_rails52?
+        # On Rails 5.2 it is no longer safe to re-use connection pools after they have been discarded
+        # This happens on webservers with forking, for example Phusion Passenger.
+        # Therefor after we clear all connections we reinitialize the shards to get fresh and not discarded ConnectionPool objects
+        proxy_config.reinitialize_shards
+      end
     end
 
     def connected?
-      @shards.any? { |_k, v| v.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?
+      should_use_slaves_for_method?(method) && 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)
+      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?
+      should_use_slaves_for_method?(method) && slave_groups.try(:[], current_slave_group).present?
     end
 
     def send_queries_to_slave_group(method, *args, &block)
-      send_queries_to_balancer(@slave_groups[current_slave_group], method, *args, &block)
+      send_queries_to_balancer(slave_groups[current_slave_group], method, *args, &block)
+    end
+
+    def current_model_replicated?
+      replicated && (current_model.try(:replicated) || fully_replicated?)
+    end
+    
+    def initialize_schema_migrations_table
+      if Octopus.atleast_rails52?
+        select_connection.transaction { ActiveRecord::SchemaMigration.create_table }
+      else 
+        select_connection.initialize_schema_migrations_table
+      end
+    end
+    
+    def initialize_metadata_table
+      select_connection.transaction { ActiveRecord::InternalMetadata.create_table }
     end
 
     protected
 
+    # @thiagopradi - This legacy method missing logic will be keep for a while for compatibility
+    # and will be removed when Octopus 1.0 will be released.
+    # We are planning to migrate to a much stable logic for the Proxy that doesn't require method missing.
+    def legacy_method_missing_logic(method, *args, &block)
+      if should_clean_connection_proxy?(method)
+        conn = select_connection
+        clean_connection_proxy
+        conn.send(method, *args, &block)
+      elsif should_send_queries_to_shard_slave_group?(method)
+        send_queries_to_shard_slave_group(method, *args, &block)
+      elsif should_send_queries_to_slave_group?(method)
+        send_queries_to_slave_group(method, *args, &block)
+      elsif should_send_queries_to_replicated_databases?(method)
+        send_queries_to_selected_slave(method, *args, &block)
+      else
+        val = select_connection.send(method, *args, &block)
+
+        if val.instance_of? ActiveRecord::Result
+          val.current_shard = shard_name
+        end
+
+        val
+      end
+    end
+
     # Ensure that a single failing slave doesn't take down the entire application
     def with_each_healthy_shard
-      @shards.each do |shard_name, v|
+      shards.each do |shard_name, v|
         begin
           yield(v)
         rescue => e
@@ -339,17 +246,10 @@ module Octopus
         end
       end
 
-      conn_handler = ActiveRecord::Base.connection_handler
-      if conn_handler.respond_to?(:connection_pool_list)
-        # Rails 4+
-        ar_pools = conn_handler.connection_pool_list
-      else
-        # Rails 3.2
-        ar_pools = conn_handler.connection_pools.values
-      end
+      ar_pools = ActiveRecord::Base.connection_handler.connection_pool_list
 
       ar_pools.each do |pool|
-        next if pool == @shards[:master] # Already handled this
+        next if pool == shards[:master] # Already handled this
 
         begin
           yield(pool)
@@ -363,57 +263,20 @@ module Octopus
       end
     end
 
-    def connection_pool_for(adapter, config)
-      if Octopus.rails4?
-        arg = ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(adapter.dup, config)
-      else
-        arg = ActiveRecord::Base::ConnectionSpecification.new(adapter.dup, config)
-      end
-
-      ActiveRecord::ConnectionAdapters::ConnectionPool.new(arg)
-    end
-
-    def initialize_adapter(adapter)
-      @adapters << adapter
-      begin
-        require "active_record/connection_adapters/#{adapter}_adapter"
-      rescue LoadError
-        raise "Please install the #{adapter} adapter: `gem install activerecord-#{adapter}-adapter` (#{$ERROR_INFO})"
-      end
-    end
-
-    def resolve_string_connection(spec)
-      if Octopus.rails41?
-        resolver = ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new({})
-        HashWithIndifferentAccess.new(resolver.spec(spec).config)
-      else
-        if Octopus.rails4?
-          resolver = ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new(spec, {})
-        else
-          resolver = ActiveRecord::Base::ConnectionSpecification::Resolver.new(spec, {})
-        end
-        HashWithIndifferentAccess.new(resolver.spec.config)
-      end
-    end
-
     def should_clean_connection_proxy?(method)
       method.to_s =~ /insert|select|execute/ && !current_model_replicated? && (!block || block != current_shard)
     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?
-    end
-
-    def current_model_replicated?
-      @replicated && (current_model.try(:replicated) || fully_replicated?)
+      replicated && method.to_s =~ /select/ && !block && !slaves_grouped?
     end
 
     def send_queries_to_selected_slave(method, *args, &block)
       if current_model.replicated || fully_replicated?
-        selected_slave = @slaves_load_balancer.next
+        selected_slave = slaves_load_balancer.next current_load_balance_options
       else
-        selected_slave = :master
+        selected_slave = Octopus.master_shard
       end
 
       send_queries_to_slave(selected_slave, method, *args, &block)
@@ -433,20 +296,24 @@ module Octopus
     end
 
     def slaves_grouped?
-      @slave_groups.present?
+      slave_groups.present?
     end
 
     # Temporarily switch `current_shard` to the next slave in a slave group and send queries to it
     # while preserving `current_shard`
     def send_queries_to_balancer(balancer, method, *args, &block)
-      send_queries_to_slave(balancer.next, method, *args, &block)
+      send_queries_to_slave(balancer.next(current_load_balance_options), method, *args, &block)
     end
 
     # Temporarily switch `current_shard` to the specified slave and send queries to it
     # while preserving `current_shard`
     def send_queries_to_slave(slave, method, *args, &block)
       using_shard(slave) do
-        select_connection.send(method, *args, &block)
+        val = select_connection.send(method, *args, &block)
+        if val.instance_of? ActiveRecord::Result
+          val.current_shard = slave
+        end
+        val
       end
     end
 
@@ -468,6 +335,8 @@ module Octopus
     # Temporarily switch `current_shard` and run the block
     def using_shard(shard, &_block)
       older_shard = current_shard
+      older_slave_group = current_slave_group
+      older_load_balance_options = current_load_balance_options
 
       begin
         unless current_model && !current_model.allowed_shard?(shard)
@@ -476,15 +345,21 @@ module Octopus
         yield
       ensure
         self.current_shard = older_shard
+        self.current_slave_group = older_slave_group
+        self.current_load_balance_options = older_load_balance_options
       end
     end
 
-    def structurally_slave?(config)
-      config.is_a?(Hash) && config.key?('adapter')
-    end
+    # Temporarily switch `current_group` and run the block
+    def using_group(group, &_block)
+      older_group = current_group
 
-    def structurally_slave_group?(config)
-      config.is_a?(Hash) && config.values.any? { |v| structurally_slave? v }
+      begin
+        self.current_group = group
+        yield
+      ensure
+        self.current_group = older_group
+      end
     end
   end
 end