switchman 3.0.2 → 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,30 +58,91 @@ 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
         end
+
+        def current_switchman_shard
+          connected_to_stack.reverse_each do |hash|
+            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.included(klass)
+      def self.prepended(klass)
         klass.singleton_class.prepend(ClassMethods)
-        klass.set_callback(:initialize, :before) do
-          @shard ||= if self.class.sharded_primary_key?
-                       Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
-                     else
-                       Shard.current(self.class.connection_classes)
-                     end
-        end
+      end
+
+      def _run_initialize_callbacks
+        @shard ||= if self.class.sharded_primary_key?
+                     Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_class_for_self))
+                   else
+                     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)
@@ -107,16 +158,16 @@ module Switchman
 
       def save(*, **)
         @shard_set_in_stone = true
-        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
+        super
       end
 
       def save!(*, **)
         @shard_set_in_stone = true
-        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
+        super
       end
 
       def destroy
-        shard.activate(self.class.connection_classes) { super }
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def clone
@@ -129,11 +180,18 @@ 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_class_for_self) do
+          db = Shard.current(self.class.connection_class_for_self).database_server
+          db.unguard { super }
+        end
+      end
+
       def hash
         self.class.sharded_primary_key? ? self.class.hash ^ global_id.hash : super
       end
@@ -149,40 +207,39 @@ module Switchman
         copy
       end
 
-      def quoted_id
-        return super unless self.class.sharded_primary_key?
-
-        # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
-        self.class.connection.quote(id)
+      def update_columns(*)
+        db = shard.database_server
+        db.unguard { super }
       end
 
-      def update_columns(*)
-        db = Shard.current(self.class.connection_classes).database_server
-        if ::GuardRail.environment == db.guard_rail_environment
-          super
+      def id_for_database
+        if self.class.sharded_primary_key?
+          # It's an int, so so it's safe to just return it without passing it through anything else
+          # In theory we should do `@attributes[@primary_key].type.serialize(id)`, but that seems to have surprising side-effects
+          id
         else
-          db.unguard { super }
+          super
         end
       end
 
       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
+              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
-          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
@@ -83,7 +83,7 @@ module Switchman
           if opts[:association]
             key_ids     = calculated_data.collect { |row| row[opts[:group_aliases].first] }
             key_records = opts[:association].klass.base_class.where(id: key_ids)
-            key_records = key_records.map { |r| [Shard.relative_id_for(r, shard, target_shard), r] }.to_h
+            key_records = key_records.to_h { |r| [Shard.relative_id_for(r, shard, target_shard), r] }
           end
 
           calculated_data.map do |row|
@@ -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,22 +1,8 @@
 # frozen_string_literal: true
 
-require 'switchman/errors'
-
 module Switchman
   module ActiveRecord
     module ConnectionPool
-      def shard
-        shard_stack.last || Shard.default
-      end
-
-      def shard_stack
-        unless (shard_stack = Thread.current.thread_variable_get(tls_key))
-          shard_stack = Concurrent::Array.new
-          Thread.current.thread_variable_set(tls_key, shard_stack)
-        end
-        shard_stack
-      end
-
       def default_schema
         connection unless @schemas
         # default shard will not switch databases immediately, so it won't be set yet
@@ -26,15 +12,15 @@ module Switchman
 
       def checkout_new_connection
         conn = super
-        conn.shard = shard
+        conn.shard = current_shard
         conn
       end
 
       def connection(switch_shard: true)
         conn = super()
-        raise NonExistentShardError if shard.new_record?
+        raise Errors::NonExistentShardError if current_shard.new_record?
 
-        switch_database(conn) if conn.shard != shard && switch_shard
+        switch_database(conn) if conn.shard != current_shard && switch_shard
         conn
       end
 
@@ -44,26 +30,18 @@ module Switchman
         flush
       end
 
-      def remove_shard!(shard)
-        synchronize do
-          # The shard might be currently active, so we need to update our own shard
-          self.shard = Shard.default if self.shard == shard
-          # Update out any connections that may be using this shard
-          @connections.each do |conn|
-            # This will also update the connection's shard to the default shard
-            switch_database(conn) if conn.shard == shard
-          end
-        end
-      end
-
       def switch_database(conn)
-        @schemas = conn.current_schemas if !@schemas && conn.adapter_name == 'PostgreSQL' && !shard.database_server.config[:shard_name]
+        @schemas = conn.current_schemas if !@schemas && conn.adapter_name == 'PostgreSQL' && !current_shard.database_server.config[:shard_name]
 
-        conn.shard = shard
+        conn.shard = current_shard
       end
 
       private
 
+      def current_shard
+        ::Rails.version < '7.0' ? connection_klass.current_switchman_shard : connection_class.current_switchman_shard
+      end
+
       def tls_key
         "#{object_id}_shard".to_sym
       end

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/migration.rb

@@ -14,16 +14,19 @@ module Switchman
 
       def connection
         conn = super
-        ::ActiveRecord::Base.connection_pool.switch_database(conn) if conn.shard != ::ActiveRecord::Base.connection_pool.shard
+        ::ActiveRecord::Base.connection_pool.switch_database(conn) if conn.shard != ::ActiveRecord::Base.current_switchman_shard
         conn
       end
     end
 
     module Migrator
-      # significant change: hash shard name, not database name
+      # significant change: just return MIGRATOR_SALT directly
+      # especially if you're going through pgbouncer, the database
+      # name you're accessing may not be consistent. it is NOT allowed
+      # to run migrations against multiple shards in the same database
+      # concurrently
       def generate_migrator_advisory_lock_id
-        shard_name_hash = Zlib.crc32(Shard.current.name)
-        ::ActiveRecord::Migrator::MIGRATOR_SALT * shard_name_hash
+        ::ActiveRecord::Migrator::MIGRATOR_SALT
       end
 
       # significant change: strip out prefer_secondary from config

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,11 +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
+        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}\""
@@ -32,7 +32,7 @@ module Switchman
       end
 
       # copy/paste; use quote_local_table_name
-      def drop_database(name) #:nodoc:
+      def drop_database(name) # :nodoc:
         execute "DROP DATABASE IF EXISTS #{quote_local_table_name(name)}"
       end
 
@@ -87,6 +87,10 @@ module Switchman
         name.quoted
       end
 
+      def with_global_table_name(&block)
+        with_local_table_name(false, &block)
+      end
+
       def with_local_table_name(enable = true) # rubocop:disable Style/OptionalBooleanParameter
         old_value = @use_local_table_name
         @use_local_table_name = enable

data/lib/switchman/active_record/predicate_builder.rb

@@ -15,7 +15,7 @@ module Switchman
         def convert_to_id(value)
           case value
           when ::ActiveRecord::Base
-            value.send(primary_key)
+            value.id
           else
             super
           end

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)
@@ -239,6 +242,10 @@ module Switchman
         connection.with_local_table_name { super }
       end
 
+      def table_name_matches?(from)
+        connection.with_global_table_name { super }
+      end
+
       def transpose_predicates(predicates,
                                source_shard,
                                target_shard,
@@ -266,11 +273,11 @@ module Switchman
           right_node = or_expr.right
           new_left_predicates = transpose_single_predicate(left_node, source_shard,
                                                            target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
-          or_expr.instance_variable_set(:@left, new_left_predicates) if new_left_predicates != left_node
           new_right_predicates = transpose_single_predicate(right_node, source_shard,
                                                             target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
-          or_expr.instance_variable_set(:@right, new_right_predicates) if new_right_predicates != right_node
-          return predicate
+          return predicate if new_left_predicates == left_node && new_right_predicates == right_node
+
+          return ::Arel::Nodes::Grouping.new ::Arel::Nodes::Or.new(new_left_predicates, new_right_predicates)
         end
         return predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
         return predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)
@@ -287,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
@@ -328,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
@@ -342,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