switchman 3.0.9 → 3.1.0

data/lib/switchman/active_record/base.rb

@@ -14,8 +14,6 @@ module Switchman
         def sharded_model
           self.abstract_class = true
 
-          return if self == UnshardedRecord
-
           Shard.send(:add_sharded_model, self)
         end
 
@@ -27,20 +25,12 @@ module Switchman
         def transaction(**)
           if self != ::ActiveRecord::Base && current_scope
             current_scope.activate do
-              db = Shard.current(connection_classes).database_server
-              if ::GuardRail.environment == db.guard_rail_environment
-                super
-              else
-                db.unguard { super }
-              end
-            end
-          else
-            db = Shard.current(connection_classes).database_server
-            if ::GuardRail.environment == db.guard_rail_environment
-              super
-            else
+              db = Shard.current(connection_class_for_self).database_server
               db.unguard { super }
             end
+          else
+            db = Shard.current(connection_class_for_self).database_server
+            db.unguard { super }
           end
         end
 
@@ -68,11 +58,56 @@ module Switchman
           end
         end
 
+        def role_overriden?(shard_id)
+          current_role(target_shard: shard_id) != current_role(without_overrides: true)
+        end
+
+        def establish_connection(config_or_env = nil)
+          raise ArgumentError, 'establish connection cannot be used on the non-current shard/role' if config_or_env.is_a?(Symbol) && config_or_env != ::Rails.env.to_sym
+
+          # Ensure we don't randomly surprise change the connection parms associated with a shard/role
+          config_or_env = nil if config_or_env == ::Rails.env.to_sym
+
+          config_or_env ||= if current_shard == ::Rails.env.to_sym && current_role == :primary
+                              :primary
+                            else
+                              "#{current_shard}/#{current_role}".to_sym
+                            end
+
+          super(config_or_env)
+        end
+
+        def connected_to_stack
+          return super if ::Rails.version < '7.0' ? Thread.current.thread_variable?(:ar_connected_to_stack) : ::ActiveSupport::IsolatedExecutionState.key?(:active_record_connected_to_stack)
+
+          ret = super
+          DatabaseServer.guard_servers
+          ret
+        end
+
+        # significant change: Allow per-shard roles
+        def current_role(without_overrides: false, target_shard: current_shard)
+          return super() if without_overrides
+
+          sharded_role = nil
+          connected_to_stack.reverse_each do |hash|
+            shard_role = hash.dig(:shard_roles, target_shard)
+            if shard_role && (hash[:klasses].include?(::ActiveRecord::Base) || hash[:klasses].include?(connection_class_for_self))
+              sharded_role = shard_role
+              break
+            end
+          end
+          # Allow a shard-specific role to be reverted to regular inheritance
+          return sharded_role if sharded_role && sharded_role != :_switchman_inherit
+
+          super()
+        end
+
         # significant change: _don't_ check if klasses.include?(Base)
         # i.e. other sharded models don't inherit the current shard of Base
         def current_shard
           connected_to_stack.reverse_each do |hash|
-            return hash[:shard] if hash[:shard] && hash[:klasses].include?(connection_classes)
+            return hash[:shard] if hash[:shard] && hash[:klasses].include?(connection_class_for_self)
           end
 
           default_shard
@@ -80,11 +115,17 @@ module Switchman
 
         def current_switchman_shard
           connected_to_stack.reverse_each do |hash|
-            return hash[:switchman_shard] if hash[:switchman_shard] && hash[:klasses].include?(connection_classes)
+            return hash[:switchman_shard] if hash[:switchman_shard] && hash[:klasses].include?(connection_class_for_self)
           end
 
           Shard.default
         end
+
+        if ::Rails.version < '7.0'
+          def connection_class_for_self
+            connection_classes
+          end
+        end
       end
 
       def self.prepended(klass)
@@ -93,15 +134,15 @@ module Switchman
 
       def _run_initialize_callbacks
         @shard ||= if self.class.sharded_primary_key?
-                     Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
+                     Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_class_for_self))
                    else
-                     Shard.current(self.class.connection_classes)
+                     Shard.current(self.class.connection_class_for_self)
                    end
         super
       end
 
       def shard
-        @shard || Shard.current(self.class.connection_classes) || Shard.default
+        @shard || Shard.current(self.class.connection_class_for_self) || Shard.default
       end
 
       def shard=(new_shard)
@@ -126,7 +167,7 @@ module Switchman
       end
 
       def destroy
-        shard.activate(self.class.connection_classes) { super }
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def clone
@@ -139,14 +180,15 @@ module Switchman
       end
 
       def transaction(**kwargs, &block)
-        shard.activate(self.class.connection_classes) do
+        shard.activate(self.class.connection_class_for_self) do
           self.class.transaction(**kwargs, &block)
         end
       end
 
       def with_transaction_returning_status
-        shard.activate(self.class.connection_classes) do
-          super
+        shard.activate(self.class.connection_class_for_self) do
+          db = Shard.current(self.class.connection_class_for_self).database_server
+          db.unguard { super }
         end
       end
 
@@ -167,9 +209,7 @@ module Switchman
 
       def update_columns(*)
         db = shard.database_server
-        return db.unguard { super } if ::GuardRail.environment != db.guard_rail_environment
-
-        super
+        db.unguard { super }
       end
 
       def id_for_database
@@ -184,22 +224,22 @@ module Switchman
 
       protected
 
-      # see also AttributeMethods#connection_classes_code_for_reflection
-      def connection_classes_for_reflection(reflection)
+      # see also AttributeMethods#connection_class_for_self_code_for_reflection
+      def connection_class_for_self_for_reflection(reflection)
         if reflection
           if reflection.options[:polymorphic]
             begin
-              read_attribute(reflection.foreign_type)&.constantize&.connection_classes || ::ActiveRecord::Base
+              read_attribute(reflection.foreign_type)&.constantize&.connection_class_for_self || ::ActiveRecord::Base
             rescue NameError
               # in case someone is abusing foreign_type to not point to an actual class
               ::ActiveRecord::Base
             end
           else
             # otherwise we can just return a symbol for the statically known type of the association
-            reflection.klass.connection_classes
+            reflection.klass.connection_class_for_self
           end
         else
-          self.class.connection_classes
+          self.class.connection_class_for_self
         end
       end
     end

data/lib/switchman/active_record/calculations.rb

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module Calculations
       def pluck(*column_names)
-        target_shard = Shard.current(klass.connection_classes)
+        target_shard = Shard.current(klass.connection_class_for_self)
         shard_count = 0
         result = activate do |relation, shard|
           shard_count += 1
@@ -109,11 +109,18 @@ module Switchman
       private
 
       def type_cast_calculated_value_switchman(value, column_name, operation)
-        type_cast_calculated_value(value, operation) do |val|
+        if ::Rails.version < '7.0'
+          type_cast_calculated_value(value, operation) do |val|
+            column = aggregate_column(column_name)
+            type ||= column.try(:type_caster) ||
+                     lookup_cast_type_from_join_dependencies(column_name.to_s) || ::ActiveRecord::Type.default_value
+            type.deserialize(val)
+          end
+        else
           column = aggregate_column(column_name)
           type ||= column.try(:type_caster) ||
-                   lookup_cast_type_from_join_dependencies(column_name.to_s) || Type.default_value
-          type.deserialize(val)
+                   lookup_cast_type_from_join_dependencies(column_name.to_s) || ::ActiveRecord::Type.default_value
+          type_cast_calculated_value(value, operation, type)
         end
       end
 

data/lib/switchman/active_record/connection_pool.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require 'switchman/errors'
-
 module Switchman
   module ActiveRecord
     module ConnectionPool
@@ -20,7 +18,7 @@ module Switchman
 
       def connection(switch_shard: true)
         conn = super()
-        raise NonExistentShardError if current_shard.new_record?
+        raise Errors::NonExistentShardError if current_shard.new_record?
 
         switch_database(conn) if conn.shard != current_shard && switch_shard
         conn
@@ -41,7 +39,7 @@ module Switchman
       private
 
       def current_shard
-        connection_klass.current_switchman_shard
+        ::Rails.version < '7.0' ? connection_klass.current_switchman_shard : connection_class.current_switchman_shard
       end
 
       def tls_key

data/lib/switchman/active_record/database_configurations.rb

@@ -3,6 +3,20 @@
 module Switchman
   module ActiveRecord
     module DatabaseConfigurations
+      # key difference: For each env name, ensure only one writable config is returned
+      # since all should point to the same data, even if multiple are writable
+      # (Picks 'primary' since it is guaranteed to exist and switchman handles activating
+      # deploy through other means)
+      def configs_for(include_replicas: false, name: nil, **)
+        res = super
+        if name && !include_replicas
+          return nil unless name.end_with?('primary')
+        elsif !include_replicas
+          return res.select { |config| config.name.end_with?('primary') }
+        end
+        res
+      end
+
       private
 
       # key difference: assumes a hybrid two-tier structure; each third tier
@@ -13,7 +27,9 @@ module Switchman
         return configs if configs.is_a?(Array)
 
         db_configs = configs.flat_map do |env_name, config|
-          roles = config.keys.select { |k| config[k].is_a?(Hash) }
+          # It would be nice to do the auto-fallback that we want here, but we haven't
+          # actually done that for years (or maybe ever) and it will be a big lift to get working
+          roles = config.keys.select { |k| config[k].is_a?(Hash) || (config[k].is_a?(Array) && config[k].all? { |ck| ck.is_a?(Hash) }) }
           base_config = config.except(*roles)
 
           name = "#{env_name}/primary"
@@ -21,7 +37,7 @@ module Switchman
           base_db = build_db_config_from_raw_config(env_name, name, base_config)
           [base_db] + roles.map do |role|
             build_db_config_from_raw_config(env_name, "#{env_name}/#{role}",
-                                            base_config.merge(config[role]).merge(replica: true))
+                                            base_config.merge(config[role].is_a?(Array) ? config[role].first : config[role]))
           end
         end
 

data/lib/switchman/active_record/finder_methods.rb

@@ -7,7 +7,7 @@ module Switchman
         return super(id) unless klass.integral_id?
 
         if shard_source_value != :implicit
-          current_shard = Shard.current(klass.connection_classes)
+          current_shard = Shard.current(klass.connection_class_for_self)
           result = activate do |relation, shard|
             current_id = Shard.relative_id_for(id, current_shard, shard)
             # current_id will be nil for non-integral id
@@ -33,7 +33,7 @@ module Switchman
       end
 
       def find_some_ordered(ids)
-        current_shard = Shard.current(klass.connection_classes)
+        current_shard = Shard.current(klass.connection_class_for_self)
         ids = ids.map { |id| Shard.relative_id_for(id, current_shard, current_shard) }
         super(ids)
       end

data/lib/switchman/active_record/model_schema.rb

@@ -6,7 +6,7 @@ module Switchman
       module ClassMethods
         def quoted_table_name
           @quoted_table_name ||= {}
-          @quoted_table_name[Shard.current(connection_classes).id] ||= connection.quote_table_name(table_name)
+          @quoted_table_name[Shard.current(connection_class_for_self).id] ||= connection.quote_table_name(table_name)
         end
       end
     end

data/lib/switchman/active_record/persistence.rb

@@ -5,18 +5,16 @@ module Switchman
     module Persistence
       # touch reads the id attribute directly, so it's not relative to the current shard
       def touch(*, **)
-        shard.activate(self.class.connection_classes) { super }
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def update_columns(*)
-        shard.activate(self.class.connection_classes) { super }
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def delete
         db = shard.database_server
-        return db.unguard { super } unless ::GuardRail.environment == db.guard_rail_environment
-
-        super
+        db.unguard { super }
       end
     end
   end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -7,7 +7,7 @@ module Switchman
       def create_database(name, options = {})
         options = { encoding: 'utf8' }.merge!(options.symbolize_keys)
 
-        option_string = options.sum do |key, value|
+        option_string = options.sum('') do |key, value|
           case key
           when :owner
             " OWNER = \"#{value}\""

data/lib/switchman/active_record/query_methods.rb

@@ -65,7 +65,7 @@ module Switchman
         when ::ActiveRecord::Relation
           Shard.default
         when nil
-          Shard.current(klass.connection_classes)
+          Shard.current(klass.connection_class_for_self)
         else
           raise ArgumentError, "invalid shard value #{shard_value}"
         end
@@ -79,7 +79,7 @@ module Switchman
         when ::ActiveRecord::Base
           shard_value.respond_to?(:associated_shards) ? shard_value.associated_shards : [shard_value.shard]
         when nil
-          [Shard.current(klass.connection_classes)]
+          [Shard.current(klass.connection_class_for_self)]
         else
           shard_value
         end
@@ -131,7 +131,7 @@ module Switchman
           id_shards = Set.new
           right.each do |value|
             local_id, id_shard = Shard.local_id_for(value)
-            id_shard ||= Shard.current(klass.connection_classes) if local_id
+            id_shard ||= Shard.current(klass.connection_class_for_self) if local_id
             id_shards << id_shard if id_shard
           end
           return if id_shards.empty?
@@ -151,10 +151,13 @@ module Switchman
           end
         when ::Arel::Nodes::BindParam
           local_id, id_shard = Shard.local_id_for(right.value.value_before_type_cast)
-          id_shard ||= Shard.current(klass.connection_classes) if local_id
+          id_shard ||= Shard.current(klass.connection_class_for_self) if local_id
+        when ::ActiveModel::Attribute
+          local_id, id_shard = Shard.local_id_for(right.value_before_type_cast)
+          id_shard ||= Shard.current(klass.connection_class_for_self) if local_id
         else
           local_id, id_shard = Shard.local_id_for(right)
-          id_shard ||= Shard.current(klass.connection_classes) if local_id
+          id_shard ||= Shard.current(klass.connection_class_for_self) if local_id
         end
 
         return if !id_shard || id_shard == primary_shard
@@ -193,9 +196,9 @@ module Switchman
           reflection = model.send(:reflection_for_integer_attribute, column)
           break if reflection
         end
-        return Shard.current(klass.connection_classes) if reflection.options[:polymorphic]
+        return Shard.current(klass.connection_class_for_self) if reflection.options[:polymorphic]
 
-        Shard.current(reflection.klass.connection_classes)
+        Shard.current(reflection.klass.connection_class_for_self)
       end
 
       def relation_and_column(attribute)
@@ -291,7 +294,7 @@ module Switchman
           if source_shard
             source_shard
           elsif type == :primary
-            Shard.current(klass.connection_classes)
+            Shard.current(klass.connection_class_for_self)
           elsif type == :foreign
             source_shard_for_foreign_key(relation, column)
           end
@@ -332,8 +335,9 @@ module Switchman
       end
 
       def transpose_predicate_value(value, current_shard, target_shard, attribute_type, remove_non_local_ids)
-        if value.is_a?(::Arel::Nodes::BindParam)
-          query_att = value.value
+        case value
+        when ::Arel::Nodes::BindParam, ::ActiveModel::Attribute
+          query_att = value.is_a?(::ActiveModel::Attribute) ? value : value.value
           current_id = query_att.value_before_type_cast
           if current_id.is_a?(::ActiveRecord::StatementCache::Substitute)
             current_id.sharded = true # mark for transposition later
@@ -346,7 +350,12 @@ module Switchman
               # make a new bind param
               value
             else
-              ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))
+              new_att = query_att.class.new(query_att.name, local_id, query_att.type)
+              if value.is_a?(::ActiveModel::Attribute)
+                new_att
+              else
+                ::Arel::Nodes::BindParam.new(new_att)
+              end
             end
           end
         else

data/lib/switchman/active_record/reflection.rb

@@ -5,7 +5,7 @@ module Switchman
     module Reflection
       module AbstractReflection
         def shard(owner)
-          if polymorphic? || klass.connection_classes == owner.class.connection_classes
+          if polymorphic? || klass.connection_class_for_self == owner.class.connection_class_for_self
             # polymorphic associations assume the same shard as the owning item
             owner.shard
           else

data/lib/switchman/active_record/relation.rb

@@ -9,13 +9,13 @@ module Switchman
 
       def initialize(*, **)
         super
-        self.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
+        self.shard_value = Shard.current(klass ? klass.connection_class_for_self : :primary) unless shard_value
         self.shard_source_value = :implicit unless shard_source_value
       end
 
       def clone
         result = super
-        result.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
+        result.shard_value = Shard.current(klass ? klass.connection_class_for_self : :primary) unless shard_value
         result
       end
 
@@ -29,35 +29,32 @@ module Switchman
       end
 
       def new(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def create(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def create!(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def to_sql
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def explain
         activate { |relation| relation.call_super(:explain, Relation) }
       end
 
-      def records
-        return @records if loaded?
-
-        results = activate { |relation| relation.call_super(:records, Relation) }
-        case shard_value
-        when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
-          @records = results
+      def load(&block)
+        if !loaded? || (::Rails.version >= '7.0' && scheduled?)
+          @records = activate { |relation| relation.send(:exec_queries, &block) }
           @loaded = true
         end
-        results
+
+        self
       end
 
       %I[update_all delete_all].each do |method|
@@ -106,19 +103,19 @@ module Switchman
       def activate(unordered: false, &block)
         shards = all_shards
         if Array === shards && shards.length == 1
-          if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_classes)
+          if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_class_for_self)
             yield(self, shards.first)
           else
-            shards.first.activate(klass.connection_classes) { yield(self, shards.first) }
+            shards.first.activate(klass.connection_class_for_self) { yield(self, shards.first) }
           end
         else
           result_count = 0
           can_order = false
-          result = Shard.with_each_shard(shards, [klass.connection_classes]) do
+          result = Shard.with_each_shard(shards, [klass.connection_class_for_self]) 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.connection_classes), :to_a)
+            relation = shard(Shard.current(klass.connection_class_for_self), :to_a)
             # do a minimal query if possible
             relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
 

data/lib/switchman/active_record/statement_cache.rb

@@ -33,12 +33,12 @@ module Switchman
           primary_value = params[primary_index]
           target_shard = Shard.local_id_for(primary_value)[1]
         end
-        current_shard = Shard.current(klass.connection_classes)
+        current_shard = Shard.current(klass.connection_class_for_self)
         target_shard ||= current_shard
 
         bind_values = bind_map.bind(params, current_shard, target_shard)
 
-        target_shard.activate(klass.connection_classes) do
+        target_shard.activate(klass.connection_class_for_self) do
           sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)
           klass.find_by_sql(sql, bind_values)
         end

data/lib/switchman/active_record/table_definition.rb

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module TableDefinition
       def column(name, type, limit: nil, **)
-        Engine.foreign_key_check(name, type, limit: limit)
+        Switchman.foreign_key_check(name, type, limit: limit)
         super
       end
     end

data/lib/switchman/active_support/cache.rb

@@ -4,6 +4,22 @@ module Switchman
   module ActiveSupport
     module Cache
       module ClassMethods
+        def lookup_stores(cache_store_config)
+          result = {}
+          cache_store_config.each do |key, value|
+            next if value.is_a?(String)
+
+            result[key] = ::ActiveSupport::Cache.lookup_store(value)
+          end
+
+          cache_store_config.each do |key, value| # rubocop:disable Style/CombinableLoops
+            next unless value.is_a?(String)
+
+            result[key] = result[value]
+          end
+          result
+        end
+
         def lookup_store(*store_options)
           store = super
           # can't use defined?, because it's a _ruby_ autoloaded constant,