switchman 3.0.2 → 3.1.0

data/lib/switchman/engine.rb

@@ -8,101 +8,27 @@ module Switchman
     config.active_record.legacy_connection_handling = false
     config.active_record.writing_role = :primary
 
-    config.autoload_once_paths << File.expand_path('app/models', config.paths.path)
+    ::GuardRail.singleton_class.prepend(GuardRail::ClassMethods)
 
-    def self.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
-
-    initializer 'switchman.initialize_cache', before: 'initialize_cache' do
-      require 'switchman/active_support/cache'
-      ::ActiveSupport::Cache.singleton_class.prepend(ActiveSupport::Cache::ClassMethods)
-
-      # if we haven't already setup our cache map out-of-band, set it up from
-      # config.cache_store now. behaves similarly to Rails' default
-      # initialize_cache initializer, but for each value in the map, rather
-      # than just Rails.cache. if config.cache_store is a flat value, uses it
-      # to fill just the Rails.env entry in the cache map.
-      unless Switchman.config[:cache_map].present?
-        cache_store_config = ::Rails.configuration.cache_store
-        cache_store_config = { ::Rails.env => cache_store_config } unless cache_store_config.is_a?(Hash)
-
-        Switchman.config[:cache_map] = Engine.lookup_stores(cache_store_config)
-      end
-
-      # if the configured cache map (either from before, or as populated from
-      # config.cache_store) didn't have an entry for Rails.env, add one using
-      # lookup_store(nil); matches the behavior of Rails' default
-      # initialize_cache initializer when config.cache_store is nil.
-      unless Switchman.config[:cache_map].key?(::Rails.env)
-        value = ::ActiveSupport::Cache.lookup_store(nil)
-        Switchman.config[:cache_map][::Rails.env] = value
-      end
-
-      middlewares = Switchman.config[:cache_map].values.map do |store|
-        store.middleware if store.respond_to?(:middleware)
-      end.compact.uniq
-      middlewares.each do |middleware|
-        config.middleware.insert_before('Rack::Runtime', middleware)
-      end
-
-      # prevent :initialize_cache from trying to (or needing to) set
-      # Rails.cache. once our switchman.extend_ar initializer (below) runs
-      # Rails.cache will be overridden to pull appropriate values from the
-      # cache map, but between now and then, Rails.cache should return the
-      # Rails.env entry in the cache map.
-      ::Rails.cache = Switchman.config[:cache_map][::Rails.env]
+    # after :initialize_dependency_mechanism to ensure autoloading is configured for any downstream initializers that care
+    # In rails 7.0 we should be able to just use an explicit after on configuring the once autoloaders and not need to go monkey around with initializer order
+    if ::Rails.version < '7.0'
+      initialize_dependency_mechanism = ::Rails::Application::Bootstrap.initializers.find { |i| i.name == :initialize_dependency_mechanism }
+      initialize_dependency_mechanism.instance_variable_get(:@options)[:after] = :set_autoload_paths
     end
 
-    initializer 'switchman.extend_ar', before: 'active_record.initialize_database' do
+    initializer 'switchman.active_record_patch',
+                before: 'active_record.initialize_database',
+                after: (::Rails.version < '7.0' ? :initialize_dependency_mechanism : :setup_once_autoloader) do
       ::ActiveSupport.on_load(:active_record) do
-        require 'switchman/active_record/abstract_adapter'
-        require 'switchman/active_record/association'
-        require 'switchman/active_record/attribute_methods'
-        require 'switchman/active_record/base'
-        require 'switchman/active_record/calculations'
-        require 'switchman/active_record/connection_pool'
-        require 'switchman/active_record/database_configurations'
-        require 'switchman/active_record/database_configurations/database_config'
-        require 'switchman/active_record/finder_methods'
-        require 'switchman/active_record/log_subscriber'
-        require 'switchman/active_record/migration'
-        require 'switchman/active_record/model_schema'
-        require 'switchman/active_record/persistence'
-        require 'switchman/active_record/predicate_builder'
-        require 'switchman/active_record/query_cache'
-        require 'switchman/active_record/query_methods'
-        require 'switchman/active_record/reflection'
-        require 'switchman/active_record/relation'
-        require 'switchman/active_record/spawn_methods'
-        require 'switchman/active_record/statement_cache'
-        require 'switchman/active_record/tasks/database_tasks'
-        require 'switchman/active_record/type_caster'
-        require 'switchman/arel'
-        require 'switchman/call_super'
-        require 'switchman/rails'
-        require 'switchman/guard_rail/relation'
-        require 'switchman/standard_error'
-
-        ::StandardError.include(StandardError)
+        # Switchman requires postgres, so just always load the pg adapter
+        require 'active_record/connection_adapters/postgresql_adapter'
 
         self.default_shard = ::Rails.env.to_sym
         self.default_role = :primary
 
-        include ActiveRecord::Base
-        include ActiveRecord::AttributeMethods
+        prepend ActiveRecord::Base
+        prepend ActiveRecord::AttributeMethods
         include ActiveRecord::Persistence
         singleton_class.prepend ActiveRecord::ModelSchema::ClassMethods
 
@@ -111,22 +37,24 @@ module Switchman
         ::ActiveRecord::StatementCache::BindMap.prepend(ActiveRecord::StatementCache::BindMap)
         ::ActiveRecord::StatementCache::Substitute.send(:attr_accessor, :primary, :sharded)
 
-        ::ActiveRecord::Associations::CollectionAssociation.prepend(ActiveRecord::CollectionAssociation)
-        ::ActiveRecord::Associations::HasOneAssociation.prepend(ActiveRecord::ForeignAssociation)
-        ::ActiveRecord::Associations::HasManyAssociation.prepend(ActiveRecord::ForeignAssociation)
+        ::ActiveRecord::Associations::CollectionAssociation.prepend(ActiveRecord::Associations::CollectionAssociation)
+        ::ActiveRecord::Associations::HasOneAssociation.prepend(ActiveRecord::Associations::ForeignAssociation)
+        ::ActiveRecord::Associations::HasManyAssociation.prepend(ActiveRecord::Associations::ForeignAssociation)
 
         ::ActiveRecord::PredicateBuilder.singleton_class.prepend(ActiveRecord::PredicateBuilder)
 
-        prepend(ActiveRecord::AutosaveAssociation)
+        prepend(ActiveRecord::Associations::AutosaveAssociation)
 
-        ::ActiveRecord::Associations::Association.prepend(ActiveRecord::Association)
-        ::ActiveRecord::Associations::BelongsToAssociation.prepend(ActiveRecord::BelongsToAssociation)
-        ::ActiveRecord::Associations::CollectionProxy.include(ActiveRecord::CollectionProxy)
+        ::ActiveRecord::Associations::Association.prepend(ActiveRecord::Associations::Association)
+        ::ActiveRecord::Associations::BelongsToAssociation.prepend(ActiveRecord::Associations::BelongsToAssociation)
+        ::ActiveRecord::Associations::CollectionProxy.include(ActiveRecord::Associations::CollectionProxy)
 
-        ::ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecord::Preloader::Association)
+        ::ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecord::Associations::Preloader::Association)
+        ::ActiveRecord::Associations::Preloader::Association::LoaderQuery.prepend(ActiveRecord::Associations::Preloader::Association::LoaderQuery) unless ::Rails.version < '7.0'
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::AbstractAdapter)
         ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ActiveRecord::ConnectionPool)
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::QueryCache)
+        ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter)
 
         ::ActiveRecord::DatabaseConfigurations.prepend(ActiveRecord::DatabaseConfigurations)
         ::ActiveRecord::DatabaseConfigurations::DatabaseConfig.prepend(ActiveRecord::DatabaseConfigurations::DatabaseConfig)
@@ -150,60 +78,74 @@ module Switchman
         ::ActiveRecord::Relation.include(CallSuper)
 
         ::ActiveRecord::PredicateBuilder::AssociationQueryValue.prepend(ActiveRecord::PredicateBuilder::AssociationQueryValue)
+        ::ActiveRecord::PredicateBuilder::PolymorphicArrayValue.prepend(ActiveRecord::PredicateBuilder::AssociationQueryValue)
 
         ::ActiveRecord::Tasks::DatabaseTasks.singleton_class.prepend(ActiveRecord::Tasks::DatabaseTasks)
 
+        ::ActiveRecord::TestFixtures.prepend(ActiveRecord::TestFixtures)
+
         ::ActiveRecord::TypeCaster::Map.include(ActiveRecord::TypeCaster::Map)
         ::ActiveRecord::TypeCaster::Connection.include(ActiveRecord::TypeCaster::Connection)
 
-        ::Rails.singleton_class.prepend(Rails::ClassMethods)
-
         ::Arel::Table.prepend(Arel::Table)
         ::Arel::Visitors::ToSql.prepend(Arel::Visitors::ToSql)
-      end
-    end
-
-    def self.foreign_key_check(name, type, limit: nil)
-      puts "WARNING: All foreign keys need to be 8-byte integers. #{name} looks like a foreign key. If so, please add the option: `:limit => 8`" if name.to_s =~ /_id\z/ && type.to_s == 'integer' && limit.to_i < 8
-    end
 
-    initializer 'switchman.extend_connection_adapters', after: 'active_record.initialize_database' do
-      ::ActiveSupport.on_load(:active_record) do
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.descendants.each do |klass|
           klass.prepend(ActiveRecord::AbstractAdapter::ForeignKeyCheck)
         end
 
-        require 'switchman/active_record/table_definition'
         ::ActiveRecord::ConnectionAdapters::TableDefinition.prepend(ActiveRecord::TableDefinition)
-
-        if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
-          require 'switchman/active_record/postgresql_adapter'
-          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter)
-        end
-
-        Shard.send(:initialize_sharding)
       end
+      # Ensure that ActiveRecord::Base is always loaded before any app-level initializers can go try to load Switchman::Shard or we get a loop
+      ::ActiveRecord::Base
     end
 
-    initializer 'switchman.eager_load' do
-      ::ActiveSupport.on_load(:before_eager_load) do
-        # This needs to be loaded before Switchman::Shard, otherwise it won't autoload it correctly
-        require 'active_record/base'
+    initializer 'switchman.error_patch', after: 'active_record.initialize_database' do
+      ::ActiveSupport.on_load(:active_record) do
+        ::StandardError.include(StandardError)
       end
     end
 
-    initializer 'switchman.extend_guard_rail', before: 'switchman.extend_ar' do
-      ::ActiveSupport.on_load(:active_record) do
-        require 'switchman/guard_rail'
+    initializer 'switchman.initialize_cache', before: :initialize_cache, after: 'active_record.initialize_database' do
+      ::ActiveSupport::Cache.singleton_class.prepend(ActiveSupport::Cache::ClassMethods)
 
-        ::GuardRail.singleton_class.prepend(GuardRail::ClassMethods)
+      # if we haven't already setup our cache map out-of-band, set it up from
+      # config.cache_store now. behaves similarly to Rails' default
+      # initialize_cache initializer, but for each value in the map, rather
+      # than just Rails.cache. if config.cache_store is a flat value, uses it
+      # to fill just the Rails.env entry in the cache map.
+      unless Switchman.config[:cache_map].present?
+        cache_store_config = ::Rails.configuration.cache_store
+        cache_store_config = { ::Rails.env => cache_store_config } unless cache_store_config.is_a?(Hash)
+
+        Switchman.config[:cache_map] = ::ActiveSupport::Cache.lookup_stores(cache_store_config)
       end
-    end
 
-    initializer 'switchman.extend_controller', after: 'guard_rail.extend_ar' do
-      ::ActiveSupport.on_load(:action_controller) do
-        require 'switchman/action_controller/caching'
+      # if the configured cache map (either from before, or as populated from
+      # config.cache_store) didn't have an entry for Rails.env, add one using
+      # lookup_store(nil); matches the behavior of Rails' default
+      # initialize_cache initializer when config.cache_store is nil.
+      unless Switchman.config[:cache_map].key?(::Rails.env)
+        value = ::ActiveSupport::Cache.lookup_store(nil)
+        Switchman.config[:cache_map][::Rails.env] = value
+      end
 
+      middlewares = Switchman.config[:cache_map].values.map do |store|
+        store.middleware if store.respond_to?(:middleware)
+      end.compact.uniq
+      middlewares.each do |middleware|
+        config.middleware.insert_before('Rack::Runtime', middleware)
+      end
+
+      # prevent :initialize_cache from trying to (or needing to) set
+      # Rails.cache. once our switchman.extend_ar initializer (below) runs
+      # Rails.cache will be overridden to pull appropriate values from the
+      # cache map, but between now and then, Rails.cache should return the
+      # Rails.env entry in the cache map.
+      ::Rails.cache = Switchman.config[:cache_map][::Rails.env]
+      ::Rails.singleton_class.prepend(Rails::ClassMethods)
+
+      ::ActiveSupport.on_load(:action_controller) do
         ::ActionController::Base.include(ActionController::Caching)
       end
     end

data/lib/switchman/errors.rb

@@ -1,7 +1,9 @@
 # frozen_string_literal: true
 
 module Switchman
-  class NonExistentShardError < RuntimeError; end
+  module Errors
+    class NonExistentShardError < RuntimeError; end
 
-  class ParallelShardExecError < RuntimeError; end
+    class ParallelShardExecError < RuntimeError; end
+  end
 end

data/lib/switchman/guard_rail/relation.rb

@@ -5,21 +5,18 @@ module Switchman
     module Relation
       def exec_queries(*args)
         if lock_value
-          db = Shard.current(connection_classes).database_server
-          return db.unguard { super } if ::GuardRail.environment != db.guard_rail_environment
+          db = Shard.current(connection_class_for_self).database_server
+          db.unguard { super }
+        else
+          super
         end
-        super
       end
 
       %w[update_all delete_all].each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
           def #{method}(*args)
-            db = Shard.current(connection_classes).database_server
-            if ::GuardRail.environment != db.guard_rail_environment
-              db.unguard { super }
-            else
-              super
-            end
+            db = Shard.current(connection_class_for_self).database_server
+            db.unguard { super }
           end
         RUBY
       end

data/lib/switchman/guard_rail.rb

@@ -3,6 +3,11 @@
 module Switchman
   module GuardRail
     module ClassMethods
+      def environment
+        # no overrides so we get the global role, not the role for the default shard
+        ::ActiveRecord::Base.current_role(without_overrides: true)
+      end
+
       def activate(role)
         DatabaseServer.send(:reference_role, role)
         super

data/lib/switchman/parallel.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require 'parallel'
+
+module Switchman
+  module Parallel
+    module UndumpableException
+      def initialize(original)
+        super
+        @active_shards = original.instance_variable_get(:@active_shards)
+        current_shard
+      end
+    end
+
+    class QuietExceptionWrapper
+      attr_accessor :name
+
+      def initialize(name, wrapper)
+        @name = name
+        @wrapper = wrapper
+      end
+
+      def exception
+        @wrapper.exception
+      end
+    end
+
+    class UndumpableResult
+      attr_reader :name
+
+      def initialize(result)
+        @name = result.inspect
+      end
+
+      def inspect
+        "#<UndumpableResult:#{name}>"
+      end
+    end
+
+    class ResultWrapper
+      attr_reader :result
+
+      def initialize(result)
+        @result =
+          begin
+            Marshal.dump(result) && result
+          rescue
+            UndumpableResult.new(result)
+          end
+      end
+    end
+
+    class PrefixingIO
+      delegate_missing_to :@original_io
+
+      def initialize(prefix, original_io)
+        @prefix = prefix
+        @original_io = original_io
+      end
+
+      def puts(*args)
+        args.flatten.each { |arg| @original_io.puts "#{@prefix}: #{arg}" }
+      end
+    end
+  end
+end
+
+::Parallel::UndumpableException.prepend(::Switchman::Parallel::UndumpableException)

data/lib/switchman/r_spec_helper.rb

@@ -70,7 +70,10 @@ module Switchman
         Shard.default(reload: true)
         puts 'Done!'
 
+        main_pid = Process.pid
         at_exit do
+          next unless main_pid == Process.pid
+
           # preserve rspec's exit status
           status = $!.is_a?(::SystemExit) ? $!.status : nil
           puts 'Tearing down sharding for all specs'
@@ -111,28 +114,12 @@ module Switchman
           Shard.default(reload: true)
           @shard1 = Shard.find(@shard1.id)
           @shard2 = Shard.find(@shard2.id)
-          shards = [@shard2]
-          shards << @shard1 unless @shard1.database_server == Shard.default.database_server
-          shards.each do |shard|
-            shard.activate do
-              ::ActiveRecord::Base.connection.begin_transaction joinable: false
-            end
-          end
         end
       end
 
       klass.after do
         next if @@sharding_failed
 
-        if use_transactional_tests
-          shards = [@shard2]
-          shards << @shard1 unless @shard1.database_server == Shard.default.database_server
-          shards.each do |shard|
-            shard.activate do
-              ::ActiveRecord::Base.connection.rollback_transaction if ::ActiveRecord::Base.connection.transaction_open?
-            end
-          end
-        end
         # clean up after specs
         DatabaseServer.all.each do |ds|
           if ds.fake? && ds != @shard2.database_server
@@ -143,7 +130,8 @@ module Switchman
       end
 
       klass.after(:all) do
-        Shard.connection.update("TRUNCATE #{Shard.quoted_table_name} CASCADE")
+        # Don't truncate because that can create some fun cross-connection lock contention
+        Shard.delete_all
         Switchman.cache.delete('default_shard')
         Shard.default(reload: true)
       end

data/lib/switchman/rails.rb

@@ -4,10 +4,7 @@ module Switchman
   module Rails
     module ClassMethods
       def self.prepended(klass)
-        # in Rails 4+, the Rails.cache= method was used during bootstrap to set
-        # Rails.cache(_without_sharding) to the value from the config file. but now
-        # that that's done (the bootstrap happened before this module is included
-        # into Rails), we want to make sure no one tries to assign to Rails.cache,
+        # we want to make sure no one tries to assign to Rails.cache,
         # because it would be wrong w.r.t. sharding.
         klass.send(:remove_method, :cache=)
       end