switchman 1.13.3 → 2.2.0

data/lib/switchman/active_record/query_methods.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module QueryMethods
@@ -227,6 +229,14 @@ module Switchman
         connection.with_local_table_name { super }
       end
 
+      def arel_column(columns)
+        connection.with_local_table_name { super }
+      end
+
+      def table_name_matches?(from)
+        connection.with_global_table_name { super }
+      end
+
       # semi-private
       public
       def transpose_predicates(predicates,
@@ -236,6 +246,18 @@ module Switchman
                                binds: nil,
                                dup_binds_on_mutation: false)
         result = predicates.map do |predicate|
+          if ::Rails.version >= '5.2' && predicate.is_a?(::Arel::Nodes::And)
+            new_predicates, _binds = transpose_predicates(predicate.children, source_shard, target_shard,
+              remove_nonlocal_primary_keys,
+              binds: binds,
+              dup_binds_on_mutation: dup_binds_on_mutation)
+            next (if new_predicates == predicate.children
+              predicate
+            else
+              ::Arel::Nodes::And.new(new_predicates)
+            end)
+          end
+
           next predicate unless predicate.is_a?(::Arel::Nodes::Binary)
           next predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)
           relation, column = relation_and_column(predicate.left)
@@ -244,7 +266,7 @@ module Switchman
           remove = true if type == :primary &&
               remove_nonlocal_primary_keys &&
               predicate.left.relation.model == klass &&
-              predicate.is_a?(::Arel::Nodes::Equality)
+              (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::In))
 
           current_source_shard =
               if source_shard
@@ -259,7 +281,7 @@ module Switchman
             new_right_value =
               case predicate.right
               when Array
-                predicate.right.map {|val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove) }
+                predicate.right.map {|val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove).presence }.compact
               else
                 transpose_predicate_value(predicate.right, current_source_shard, target_shard, type, remove)
               end
@@ -334,7 +356,7 @@ module Switchman
             value
           else
             local_id = Shard.relative_id_for(current_id, current_shard, target_shard) || current_id
-            local_id = [] if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
+            return nil if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
             if current_id != local_id
               # make a new bind param
               ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))

data/lib/switchman/active_record/reflection.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Reflection
@@ -26,15 +28,26 @@ module Switchman
         # this technically belongs on AssociationReflection, but we put it on
         # ThroughReflection as well, instead of delegating to its internal
         # HasManyAssociation, losing its proper `klass`
-        def association_scope_cache(conn, owner, &block)
-          key = conn.prepared_statements
-          if polymorphic?
-            key = [key, owner._read_attribute(@foreign_type)]
+        if ::Rails.version < '6.0.4'
+          def association_scope_cache(conn, owner, &block)
+            key = conn.prepared_statements
+            if polymorphic?
+              key = [key, owner._read_attribute(@foreign_type)]
+            end
+            key = [key, shard(owner).id].flatten
+            @association_scope_cache[key] ||= @scope_lock.synchronize {
+              @association_scope_cache[key] ||= (::Rails.version >= "5.2" ? ::ActiveRecord::StatementCache.create(conn, &block) : block.call)
+            }
+          end
+        else
+          def association_scope_cache(klass, owner, &block)
+            key = self
+            if polymorphic?
+              key = [key, owner._read_attribute(@foreign_type)]
+            end
+            key = [key, shard(owner).id].flatten
+            klass.cached_find_by_statement(key, &block)
           end
-          key = [key, shard(owner).id].flatten
-          @association_scope_cache[key] ||= @scope_lock.synchronize {
-            @association_scope_cache[key] ||= (::Rails.version >= "5.2" ? ::ActiveRecord::StatementCache.create(conn, &block) : block.call)
-          }
         end
       end
 

data/lib/switchman/active_record/relation.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Relation
@@ -5,7 +7,7 @@ module Switchman
         klass::SINGLE_VALUE_METHODS.concat [ :shard, :shard_source ]
       end
 
-      def initialize(*args)
+      def initialize(*, **)
         super
         self.shard_value = Shard.current(klass ? klass.shard_category : :primary) unless shard_value
         self.shard_source_value = :implicit unless shard_source_value
@@ -17,7 +19,7 @@ module Switchman
         result
       end
 
-      def merge(*args)
+      def merge(*)
         relation = super
         if relation.shard_value != self.shard_value && relation.shard_source_value == :implicit
           relation.shard_value = self.shard_value
@@ -26,15 +28,15 @@ module Switchman
         relation
       end
 
-      def new(*args, &block)
+      def new(*, &block)
         primary_shard.activate(klass.shard_category) { super }
       end
 
-      def create(*args, &block)
+      def create(*, &block)
         primary_shard.activate(klass.shard_category) { super }
       end
 
-      def create!(*args, &block)
+      def create!(*, &block)
         primary_shard.activate(klass.shard_category) { super }
       end
 
@@ -60,7 +62,7 @@ module Switchman
       %I{update_all delete_all}.each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
           def #{method}(*args)
-            result = self.activate { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
+            result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
             result = result.sum if result.is_a?(Array)
             result
           end
@@ -95,7 +97,7 @@ module Switchman
         end
       end
 
-      def activate(&block)
+      def activate(unordered: false, &block)
         shards = all_shards
         if (Array === shards && shards.length == 1)
           if shards.first == DefaultShard || shards.first == Shard.current(klass.shard_category)
@@ -104,10 +106,64 @@ module Switchman
             shards.first.activate(klass.shard_category) { yield(self, shards.first) }
           end
         else
-          # TODO: implement local limit to avoid querying extra shards
-          Shard.with_each_shard(shards, [klass.shard_category]) do
-            shard(Shard.current(klass.shard_category), :to_a).activate(&block)
+          result_count = 0
+          can_order = false
+          result = Shard.with_each_shard(shards, [klass.shard_category]) do
+            # don't even query other shards if we're already past the limit
+            next if limit_value && result_count >= limit_value && order_values.empty?
+
+            relation = shard(Shard.current(klass.shard_category), :to_a)
+            # do a minimal query if possible
+            relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
+
+            shard_results = relation.activate(&block)
+
+            if shard_results.present? && !unordered
+              can_order ||= can_order_cross_shard_results? unless order_values.empty?
+              raise OrderOnMultiShardQuery if !can_order && !order_values.empty? && result_count.positive?
+
+              result_count += shard_results.is_a?(Array) ? shard_results.length : 1
+            end
+            shard_results
+          end
+
+          result = reorder_cross_shard_results(result) if can_order
+          result.slice!(limit_value..-1) if limit_value
+          result
+        end
+      end
+
+      def can_order_cross_shard_results?
+        # we only presume to be able to post-sort the most basic of orderings
+        order_values.all? { |ov| ov.is_a?(::Arel::Nodes::Ordering) && ov.expr.is_a?(::Arel::Attributes::Attribute) }
+      end
+
+      def reorder_cross_shard_results(results)
+        results.sort! do |l, r|
+          result = 0
+          order_values.each do |ov|
+            if !l.is_a?(::ActiveRecord::Base)
+              a, b = l, r
+            elsif l.respond_to?(ov.expr.name)
+              a = l.send(ov.expr.name)
+              b = r.send(ov.expr.name)
+            else
+              a = l.attributes[ov.expr.name]
+              b = r.attributes[ov.expr.name]
+            end
+            next if a == b
+
+            if a.nil? || b.nil?
+              result = 1 if a.nil?
+              result *= -1 if ov.is_a?(::Arel::Nodes::Descending)
+            else
+              result = a <=> b
+            end
+
+            result *= -1 if ov.is_a?(::Arel::Nodes::Descending)
+            break unless result.zero?
           end
+          result
         end
       end
     end

data/lib/switchman/active_record/spawn_methods.rb

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

data/lib/switchman/active_record/statement_cache.rb

@@ -1,8 +1,10 @@
+# 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
 
           if ::Rails.version >= "5.2"
@@ -46,37 +48,16 @@ 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
 
-      if ::Rails.version < '5.1'
-        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)
-        end
-      elsif ::Rails.version < '5.2'
-        def generic_query_builder(connection)
-          @query_builder ||= connection.cacheable_query(self.class, @arel)
-        end
-
+      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 generic_query_builder(connection)
-          @query_builder ||= connection.cacheable_query(self.class, @arel).first
-        end
-
         def qualified_query_builder(shard, klass)
           @qualified_query_builders[shard.id] ||= klass.connection.cacheable_query(self.class, @arel).first
         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

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,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module Arel
     module Table

data/lib/switchman/call_super.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module CallSuper
     def super_method_above(method_name, above_module)

data/lib/switchman/connection_pool_proxy.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require 'switchman/schema_cache'
 
 module Switchman
@@ -23,20 +25,27 @@ 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
@@ -46,21 +55,21 @@ module Switchman
       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,6 +80,14 @@ module Switchman
       end
     end
 
+    def get_schema_cache(_connection)
+      @schema_cache
+    end
+
+    def set_schema_cache(cache)
+      @schema_cache.copy_values(cache)
+    end
+
     %w{release_connection
        disconnect!
        flush!
@@ -98,6 +115,10 @@ module Switchman
       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
@@ -105,7 +126,7 @@ module Switchman
     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
 
@@ -113,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
@@ -123,27 +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
-      args = [config, "#{config[:adapter]}_connection"]
-      args.unshift(pool_key.join("/"))
-      spec = ::ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(*args)
+      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
-        if ::Rails.version >= '5.0.1'
-          pool.enable_query_cache! if !@connection_pools.empty? && @connection_pools.first.last.query_cache_enabled
-        end
+        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

data/lib/switchman/database_server.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "securerandom"
 
 module Switchman
@@ -40,8 +42,14 @@ module Switchman
       def database_servers
         unless @database_servers
           @database_servers = {}.with_indifferent_access
-          ::ActiveRecord::Base.configurations.each do |(id, config)|
-            @database_servers[id] = DatabaseServer.new(id, config)
+          if ::Rails.version >= '6.0'
+            ::ActiveRecord::Base.configurations.configurations.each do |config|
+              @database_servers[config.env_name] = DatabaseServer.new(config.env_name, config.config)
+            end
+          else
+            ::ActiveRecord::Base.configurations.each do |(id, config)|
+              @database_servers[id] = DatabaseServer.new(id, config)
+            end
           end
         end
         @database_servers
@@ -63,12 +71,12 @@ module Switchman
       @fake
     end
 
-    def config(environment = :master)
+    def config(environment = :primary)
       @configs[environment] ||= begin
         if @config[environment].is_a?(Array)
           @config[environment].map do |config|
             config = @config.merge((config || {}).symbolize_keys)
-            # make sure Shackles doesn't get any brilliant ideas about choosing the first possible server
+            # make sure GuardRail doesn't get any brilliant ideas about choosing the first possible server
             config.delete(environment)
             config
           end
@@ -80,33 +88,33 @@ module Switchman
       end
     end
 
-    def shackles_environment
-      @shackles_environment || ::Shackles.environment
+    def guard_rail_environment
+      @guard_rail_environment || ::GuardRail.environment
     end
 
     # locks this db to a specific environment, except for
     # when doing writes (then it falls back to the current
-    # value of Shackles.environment)
-    def shackle!(environment = :slave)
-      @shackles_environment = environment
+    # value of GuardRail.environment)
+    def guard!(environment = :secondary)
+      @guard_rail_environment = environment
     end
 
-    def unshackle!
-      @shackles_environment = nil
+    def unguard!
+      @guard_rail_environment = nil
     end
 
-    def unshackle
-      old_env = @shackles_environment
-      unshackle!
+    def unguard
+      old_env = @guard_rail_environment
+      unguard!
       yield
     ensure
-      shackle!(old_env)
+      guard!(old_env)
     end
 
     def shareable?
       @shareable_environment_key ||= []
-      environment = shackles_environment
-      explicit_user = ::Shackles.global_config[:username]
+      environment = guard_rail_environment
+      explicit_user = ::GuardRail.global_config[:username]
       return @shareable if @shareable_environment_key == [environment, explicit_user]
       @shareable_environment_key = [environment, explicit_user]
       if explicit_user
@@ -179,13 +187,18 @@ module Switchman
         shard = Shard.create!(:id => shard_id,
                               :name => name,
                               :database_server => self)
+        schema_already_existed = false
 
         begin
           self.class.creating_new_shard = true
           shard.activate(*Shard.categories) do
-            ::Shackles.activate(:deploy) do
+            ::GuardRail.activate(:deploy) do
               begin
                 if create_statement
+                  if (::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}"))
+                    schema_already_existed = true
+                    raise "This schema already exists; cannot overwrite"
+                  end
                   Array(create_statement.call).each do |stmt|
                     ::ActiveRecord::Base.connection.execute(stmt)
                   end
@@ -223,7 +236,9 @@ module Switchman
           shard
         rescue
           shard.destroy
-          shard.drop_database rescue nil
+          unless schema_already_existed
+            shard.drop_database rescue nil
+          end
           reset_column_information unless create_schema == false rescue nil
           raise
         ensure

data/lib/switchman/default_shard.rb

@@ -1,8 +1,11 @@
+# frozen_string_literal: true
+
 require 'switchman/database_server'
 
 module Switchman
   class DefaultShard
     def id; 'default'; end
+    alias cache_key id
     def activate(*categories); yield; end
     def activate!(*categories); end
     def default?; true; end