switchman 1.5.21 → 2.1.5

data/lib/switchman/active_record/statement_cache.rb

@@ -1,19 +1,27 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module StatementCache
       module ClassMethods
-        def create(connection, block = Proc.new)
+        def create(connection, &block)
           relation = block.call ::ActiveRecord::StatementCache::Params.new
 
-          binds = ::Rails.version >= '5' ? relation.bound_attributes : relation.bind_values
-          bind_map = ::ActiveRecord::StatementCache::BindMap.new(binds)
-          new relation.arel, bind_map
+          if ::Rails.version >= "5.2"
+            query_builder, binds = connection.cacheable_query(self, relation.arel)
+            bind_map = ::ActiveRecord::StatementCache::BindMap.new(binds)
+            new(relation.arel, bind_map, relation.klass)
+          else
+            bind_map = ::ActiveRecord::StatementCache::BindMap.new(relation.bound_attributes)
+            new relation.arel, bind_map
+          end
         end
       end
 
-      def initialize(arel, bind_map)
+      def initialize(arel, bind_map, klass=nil)
         @arel = arel
         @bind_map = bind_map
+        @klass = klass
         @qualified_query_builders = {}
       end
 
@@ -22,7 +30,13 @@ module Switchman
       # we can make some assumptions about the shard source
       # (e.g. infer from the primary key or use the current shard)
 
-      def execute(params, klass, connection)
+      def execute(*args)
+        if ::Rails.version >= '5.2'
+          params, connection = args
+          klass = @klass
+        else
+          params, klass, connection = args
+        end
         target_shard = nil
         if primary_index = bind_map.primary_value_index
           primary_value = params[primary_index]
@@ -34,65 +48,43 @@ module Switchman
         bind_values = bind_map.bind(params, current_shard, target_shard)
 
         target_shard.activate(klass.shard_category) do
-          if connection.use_qualified_names?
-            sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)
-            klass.find_by_sql(sql, bind_values)
-          else
-            sql = generic_query_builder(connection).sql_for(bind_values, connection)
-            klass.find_by_sql(sql, bind_values)
-          end
+          sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)
+          klass.find_by_sql(sql, bind_values)
         end
       end
 
-      def generic_query_builder(connection)
-        @query_builder ||= connection.cacheable_query(@arel)
-      end
-
-      def qualified_query_builder(shard, klass)
-        @qualified_query_builders[shard.id] ||= klass.connection.cacheable_query(@arel)
+      if ::Rails.version < '5.2'
+        def qualified_query_builder(shard, klass)
+          @qualified_query_builders[shard.id] ||= klass.connection.cacheable_query(self.class, @arel)
+        end
+      else
+        def qualified_query_builder(shard, klass)
+          @qualified_query_builders[shard.id] ||= klass.connection.cacheable_query(self.class, @arel).first
+        end
       end
 
       module BindMap
         # performs id transposition here instead of query_methods.rb
         def bind(values, current_shard, target_shard)
-          if ::Rails.version >= '5'
-            bas = @bound_attributes.dup
-            @indexes.each_with_index do |offset,i|
-              ba = bas[offset]
-              if ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.sharded
-                new_value = Shard.relative_id_for(values[i], current_shard, target_shard || current_shard)
-              else
-                new_value = values[i]
-              end
-              bas[offset] = ba.with_cast_value(new_value)
-            end
-            bas
-          else
-            bvs = @bind_values.map { |pair| pair.dup }
-            @indexes.each_with_index do |offset,i|
-              bv = bvs[offset]
-              if bv[1].sharded
-                bv[1] = Shard.relative_id_for(values[i], current_shard, target_shard || current_shard)
-              else
-                bv[1] = values[i]
-              end
+          bas = @bound_attributes.dup
+          @indexes.each_with_index do |offset,i|
+            ba = bas[offset]
+            if ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.sharded
+              new_value = Shard.relative_id_for(values[i], current_shard, target_shard || current_shard)
+            else
+              new_value = values[i]
             end
-            bvs
+            bas[offset] = ba.with_cast_value(new_value)
           end
+          bas
         end
 
         def primary_value_index
-          if ::Rails.version >= '5'
-            primary_ba_index = @bound_attributes.index do |ba|
-              ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.primary
-            end
-            if primary_ba_index
-              @indexes.index(primary_ba_index)
-            end
-          else
-            if primary_bv_index = @bind_values.index{|col, sub| sub.primary}
-              @indexes.index(primary_bv_index)
-            end
+          primary_ba_index = @bound_attributes.index do |ba|
+            ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.primary
+          end
+          if primary_ba_index
+            @indexes.index(primary_ba_index)
           end
         end
       end

data/lib/switchman/active_record/table_definition.rb

@@ -1,8 +1,10 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module TableDefinition
-      def column(name, type, options = {})
-        Engine.foreign_key_check(name, type, options)
+      def column(name, type, limit: nil, **)
+        Engine.foreign_key_check(name, type, limit: limit)
         super
       end
     end

data/lib/switchman/active_record/type_caster.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module TypeCaster

data/lib/switchman/active_record/where_clause_factory.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module WhereClauseFactory
@@ -19,9 +21,10 @@ module Switchman
           super
         when Hash, ::Arel::Nodes::Node
           where_clause = super
+          binds = ::Rails.version >= "5.2" ? nil : where_clause.binds
           predicates = where_clause.send(:predicates)
-          @scope.send(:infer_shards_from_primary_key, predicates, where_clause.binds) if @scope.shard_source_value == :implicit && @scope.shard_value.is_a?(Shard)
-          predicates = @scope.transpose_predicates(predicates, nil, @scope.primary_shard, false, where_clause.binds) if @scope.shard_source_value != :explicit
+          @scope.send(:infer_shards_from_primary_key, predicates, binds) if @scope.shard_source_value == :implicit && @scope.shard_value.is_a?(Shard)
+          predicates, _new_binds = @scope.transpose_predicates(predicates, nil, @scope.primary_shard, false, binds: binds)
           where_clause.instance_variable_set(:@predicates, predicates)
           where_clause
         else

data/lib/switchman/active_support/cache.rb

@@ -1,13 +1,31 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveSupport
     module Cache
       module ClassMethods
         def lookup_store(*store_options)
           store = super
+          # can't use defined?, because it's a _ruby_ autoloaded constant,
+          # so just checking that will cause it to get required
+          if store.class.name == "ActiveSupport::Cache::RedisCacheStore" && !::ActiveSupport::Cache::RedisCacheStore.ancestors.include?(RedisCacheStore)
+            ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore)
+          end
           store.options[:namespace] ||= lambda { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
           store
         end
       end
+
+      module RedisCacheStore
+        def clear(namespace: nil, **)
+          # RedisCacheStore tries to be smart and only clear the cache under your namespace, if you have one set
+          # unfortunately, it uses the keys command, which is extraordinarily inefficient in a large redis instance
+          # fortunately, we can assume we control the entire instance, because we set up the namespacing, so just
+          # always unset it temporarily for clear calls
+          namespace = nil
+          super
+        end
+      end
     end
   end
 end

data/lib/switchman/arel.rb

@@ -1,35 +1,26 @@
+# frozen_string_literal: true
+
 module Switchman
   module Arel
     module Table
       def model
-        if ::Rails.version >= '5'
-          type_caster.model
-        else
-          engine
-        end
+        type_caster.model
       end
     end
     module Visitors
       module ToSql
         def visit_Arel_Nodes_TableAlias *args
-          if ::Rails.version < '4.2'
-            o = args.shift
-            "#{visit o.relation, *args} #{quote_local_table_name o.name}"
-          else
-            o, collector = args
-            collector = visit o.relation, collector
-            collector << " "
-            collector << quote_local_table_name(o.name)
-          end
+          o, collector = args
+          collector = visit o.relation, collector
+          collector << " "
+          collector << quote_local_table_name(o.name)
         end
 
         def visit_Arel_Attributes_Attribute *args
           o = args.first
           join_name = o.relation.table_alias || o.relation.name
           result = "#{quote_local_table_name join_name}.#{quote_column_name o.name}"
-          unless ::Rails.version < '4.2'.freeze
-            result = args.last << result
-          end
+          result = args.last << result
           result
         end
 
@@ -38,14 +29,6 @@ module Switchman
           @connection.quote_local_table_name(name)
         end
       end
-
-      module PostgreSQL
-        # the only difference is to remove caching (which only applies to Arel < 6.0/AR < 4.2)
-        def quote_table_name name
-          return name if ::Arel::Nodes::SqlLiteral === name
-          @connection.quote_table_name(name)
-        end
-      end
     end
   end
 end

data/lib/switchman/call_super.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Switchman
+  module CallSuper
+    def super_method_above(method_name, above_module)
+      method = method(method_name)
+      last_owner = method.owner
+      while method.owner != above_module
+        method = method.super_method
+        raise "Could not find super method ``#{method_name}' for #{self.class}" if method.owner == last_owner
+      end
+      method.super_method
+    end
+
+    def call_super(method, above_module, *args, &block)
+      super_method_above(method, above_module).call(*args, &block)
+    end
+  end
+end

data/lib/switchman/connection_pool_proxy.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require 'switchman/schema_cache'
 
 module Switchman
@@ -10,7 +12,7 @@ module Switchman
   end
 
   class ConnectionPoolProxy
-    delegate :spec, :connected?, :default_schema, :with_connection,
+    delegate :spec, :connected?, :default_schema, :with_connection, :query_cache_enabled, :active_connection?,
              :to => :current_pool
 
     attr_reader :category, :schema_cache
@@ -23,44 +25,51 @@ module Switchman
       @category = category
       @default_pool = default_pool
       @connection_pools = shard_connection_pools
-      @schema_cache = SchemaCache.new(self)
+      @schema_cache = default_pool.get_schema_cache(nil) if ::Rails.version >= '6'
+      @schema_cache = SchemaCache.new(self) unless @schema_cache.is_a?(SchemaCache)
+      if ::Rails.version >= '6'
+        @default_pool.set_schema_cache(@schema_cache)
+        @connection_pools.each_value do |pool|
+          pool.set_schema_cache(@schema_cache)
+        end
+      end
     end
 
     def active_shard
       Shard.current(@category)
     end
 
-    def active_shackles_environment
-      ::Rails.env.test? ? :master : active_shard.database_server.shackles_environment
+    def active_guard_rail_environment
+      ::Rails.env.test? ? :primary : active_shard.database_server.guard_rail_environment
     end
 
     def current_pool
       current_active_shard = active_shard
-      pool = self.default_pool if current_active_shard.database_server == Shard.default.database_server && active_shackles_environment == :master && (current_active_shard.default? || current_active_shard.database_server.shareable?)
+      pool = self.default_pool if current_active_shard.database_server == Shard.default.database_server && active_guard_rail_environment == :primary && (current_active_shard.default? || current_active_shard.database_server.shareable?)
       pool = @connection_pools[pool_key] ||= create_pool unless pool
       pool.shard = current_active_shard
       pool
     end
 
     def connections
-      @connection_pools.values.map(&:connections).inject([], &:+)
+      connection_pools.map(&:connections).inject([], &:+)
     end
 
-    def connection
+    def connection(switch_shard: true)
       pool = current_pool
       begin
-        connection = pool.connection
-        connection.instance_variable_set(:@schema_cache, @schema_cache)
+        connection = pool.connection(switch_shard: switch_shard)
+        connection.instance_variable_set(:@schema_cache, @schema_cache) unless ::Rails.version >= '6'
         connection
       rescue ConnectionError
-        raise if active_shard.database_server == Shard.default.database_server && active_shackles_environment == :master
-        configs = active_shard.database_server.config(active_shackles_environment)
+        raise if active_shard.database_server == Shard.default.database_server && active_guard_rail_environment == :primary
+        configs = active_shard.database_server.config(active_guard_rail_environment)
         raise unless configs.is_a?(Array)
         configs.each_with_index do |config, idx|
           pool = create_pool(config.dup)
           begin
             connection = pool.connection
-            connection.instance_variable_set(:@schema_cache, @schema_cache)
+            connection.instance_variable_set(:@schema_cache, @schema_cache) unless ::Rails.version >= '6'
           rescue ConnectionError
             raise if idx == configs.length - 1
             next
@@ -71,22 +80,53 @@ module Switchman
       end
     end
 
-    %w{release_connection disconnect! clear_reloadable_connections! verify_active_connections! clear_stale_cached_connections!}.each do |method|
-      class_eval(<<-EOS)
+    def get_schema_cache(_connection)
+      @schema_cache
+    end
+
+    def set_schema_cache(cache)
+      @schema_cache.copy_values(cache)
+    end
+
+    %w{release_connection
+       disconnect!
+       flush!
+       clear_reloadable_connections!
+       verify_active_connections!
+       clear_stale_cached_connections!
+       enable_query_cache!
+       disable_query_cache! }.each do |method|
+      class_eval(<<-RUBY, __FILE__, __LINE__ + 1)
           def #{method}
-            @connection_pools.values.each(&:#{method})
+            connection_pools.each(&:#{method})
           end
-      EOS
+      RUBY
+    end
+
+    def discard!
+      # this breaks everything if i try to pass it onto the pools and i'm not sure why
+    end
+
+    def automatic_reconnect=(value)
+      connection_pools.each { |pool| pool.automatic_reconnect = value }
     end
 
     def clear_idle_connections!(since_when)
-      @connection_pools.values.each { |pool| pool.clear_idle_connections!(since_when) }
+      connection_pools.each { |pool| pool.clear_idle_connections!(since_when) }
+    end
+
+    def remove_shard!(shard)
+      connection_pools.each { |pool| pool.remove_shard!(shard) }
     end
 
     protected
 
+    def connection_pools
+      (@connection_pools.values + [default_pool]).uniq
+    end
+
     def pool_key
-      [active_shackles_environment,
+      [active_guard_rail_environment,
         active_shard.database_server.shareable? ? active_shard.database_server.pool_key : active_shard]
     end
 
@@ -94,7 +134,7 @@ module Switchman
       shard = active_shard
       unless config
         if shard != Shard.default
-          config = shard.database_server.config(active_shackles_environment)
+          config = shard.database_server.config(active_guard_rail_environment)
           config = config.first if config.is_a?(Array)
           config = config.dup
         else
@@ -104,22 +144,28 @@ module Switchman
           # different models could be using different configs on the default
           # shard, and database server wouldn't know about that
           config = default_pool.spec.instance_variable_get(:@config)
-          if config[active_shackles_environment].is_a?(Hash)
-            config = config.merge(config[active_shackles_environment])
-          elsif config[active_shackles_environment].is_a?(Array)
-            config = config.merge(config[active_shackles_environment].first)
+          if config[active_guard_rail_environment].is_a?(Hash)
+            config = config.merge(config[active_guard_rail_environment])
+          elsif config[active_guard_rail_environment].is_a?(Array)
+            config = config.merge(config[active_guard_rail_environment].first)
           else
             config = config.dup
           end
         end
       end
-      spec = ::ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(config, "#{config[:adapter]}_connection")
+      spec = ::ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(
+        default_pool.spec.name,
+        config,
+        "#{config[:adapter]}_connection"
+      )
       # unfortunately the AR code that does this require logic can't really be
       # called in isolation
       require "active_record/connection_adapters/#{config[:adapter]}_adapter"
 
       ::ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec).tap do |pool|
         pool.shard = shard
+        pool.set_schema_cache(@schema_cache) if ::Rails.version >= '6'
+        pool.enable_query_cache! if !@connection_pools.empty? && @connection_pools.first.last.query_cache_enabled
       end
     end
   end