switchman 3.2.1 → 4.0.0

data/lib/switchman/arel.rb

@@ -13,38 +13,54 @@ module Switchman
         # rubocop:disable Naming/MethodName
         # rubocop:disable Naming/MethodParameterName
 
+        def visit_Arel_Nodes_Cte(o, collector)
+          collector << quote_local_table_name(o.name)
+          collector << " AS "
+
+          case o.materialized
+          when true
+            collector << "MATERIALIZED "
+          when false
+            collector << "NOT MATERIALIZED "
+          end
+
+          visit o.relation, collector
+        end
+
         def visit_Arel_Nodes_TableAlias(o, collector)
           collector = visit o.relation, collector
-          collector << ' '
+          collector << " "
           collector << quote_local_table_name(o.name)
         end
 
         def visit_Arel_Attributes_Attribute(o, collector)
           join_name = o.relation.table_alias || o.relation.name
-          collector << quote_local_table_name(join_name) << '.' << quote_column_name(o.name)
+          collector << quote_local_table_name(join_name) << "." << quote_column_name(o.name)
         end
 
-        def visit_Arel_Nodes_HomogeneousIn(o, collector)
-          collector.preparable = false
+        if ::Rails.version < "7.1"
+          def visit_Arel_Nodes_HomogeneousIn(o, collector)
+            collector.preparable = false
 
-          collector << quote_local_table_name(o.table_name) << '.' << quote_column_name(o.column_name)
+            collector << quote_local_table_name(o.table_name) << "." << quote_column_name(o.column_name)
 
-          collector << if o.type == :in
-                         ' IN ('
-                       else
-                         ' NOT IN ('
-                       end
+            collector << if o.type == :in
+                           " IN ("
+                         else
+                           " NOT IN ("
+                         end
 
-          values = o.casted_values
+            values = o.casted_values
 
-          if values.empty?
-            collector << @connection.quote(nil)
-          else
-            collector.add_binds(values, o.proc_for_binds, &bind_block)
-          end
+            if values.empty?
+              collector << @connection.quote(nil)
+            else
+              collector.add_binds(values, o.proc_for_binds, &bind_block)
+            end
 
-          collector << ')'
-          collector
+            collector << ")"
+            collector
+          end
         end
 
         # rubocop:enable Naming/MethodName

data/lib/switchman/call_super.rb

@@ -12,8 +12,8 @@ module Switchman
       method.super_method
     end
 
-    def call_super(method, above_module, *args, &block)
-      super_method_above(method, above_module).call(*args, &block)
+    def call_super(method, above_module, ...)
+      super_method_above(method, above_module).call(...)
     end
   end
 end

data/lib/switchman/database_server.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'securerandom'
+require "securerandom"
 
 module Switchman
   class DatabaseServer
@@ -10,19 +10,23 @@ module Switchman
       attr_accessor :creating_new_shard
       attr_reader :all_roles
 
+      include Enumerable
+
+      delegate :each, to: :all
+
       def all
         database_servers.values
       end
 
       def find(id_or_all)
         return all if id_or_all == :all
-        return id_or_all.map { |id| database_servers[id || ::Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
+        return id_or_all.filter_map { |id| database_servers[id || ::Rails.env] }.uniq if id_or_all.is_a?(Array)
 
         database_servers[id_or_all || ::Rails.env]
       end
 
       def create(settings = {})
-        raise 'database servers should be set up in database.yml' unless ::Rails.env.test?
+        raise "database servers should be set up in database.yml" unless ::Rails.env.test?
 
         id = settings[:id]
         unless id
@@ -50,6 +54,10 @@ module Switchman
         all.each { |db| db.guard! if db.config[:prefer_secondary] }
       end
 
+      def regions
+        @regions ||= all.filter_map(&:region).uniq.sort
+      end
+
       private
 
       def reference_role(role)
@@ -64,8 +72,8 @@ module Switchman
           @database_servers = {}.with_indifferent_access
           roles = []
           ::ActiveRecord::Base.configurations.configurations.each do |config|
-            if config.name.include?('/')
-              name, role = config.name.split('/')
+            if config.name.include?("/")
+              name, role = config.name.split("/")
             else
               name, role = config.env_name, config.name
             end
@@ -81,6 +89,8 @@ module Switchman
           # Do this after so that all database servers for all roles are established and we won't prematurely
           # configure a connection for the wrong role
           @all_roles = roles.uniq
+          return @database_servers if @database_servers.empty?
+
           Shard.send(:configure_connects_to)
         end
         @database_servers
@@ -110,8 +120,9 @@ module Switchman
       self.class.send(:database_servers).delete(id) if id
       Shard.sharded_models.each do |klass|
         self.class.all_roles.each do |role|
-          klass.connection_handler.remove_connection_pool(klass.connection_specification_name, role: role,
-                                                                                               shard: id.to_sym)
+          klass.connection_handler.remove_connection_pool(klass.connection_specification_name,
+                                                          role: role,
+                                                          shard: id.to_sym)
         end
       end
     end
@@ -137,16 +148,41 @@ module Switchman
         end
     end
 
+    def region
+      config[:region]
+    end
+
+    # @param region [String, Array<String>] the region(s) to check against
+    # @return true if the database server doesn't have a region, or it
+    #   matches the specified region
+    def in_region?(region)
+      !self.region || (region.is_a?(Array) ? region.include?(self.region) : self.region == region)
+    end
+
+    # @return true if the database server doesn't have a region, Switchman is
+    #   not configured with a region, or the database server's region matches
+    #   Switchman's current region
+    def in_current_region?
+      unless instance_variable_defined?(:@in_current_region)
+        @in_current_region = !region ||
+                             !Switchman.region ||
+                             region == Switchman.region
+      end
+      @in_current_region
+    end
+
     # locks this db to a specific environment, except for
     # when doing writes (then it falls back to the current
     # value of GuardRail.environment)
     def guard!(environment = :secondary)
       DatabaseServer.send(:reference_role, environment)
-      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => environment }, klasses: [::ActiveRecord::Base] }
+      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => environment },
+                                                   klasses: [::ActiveRecord::Base] }
     end
 
     def unguard!
-      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => :_switchman_inherit }, klasses: [::ActiveRecord::Base] }
+      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => :_switchman_inherit },
+                                                   klasses: [::ActiveRecord::Base] }
     end
 
     def unguard
@@ -162,7 +198,7 @@ module Switchman
 
     def shards
       if id == ::Rails.env
-        Shard.where('database_server_id IS NULL OR database_server_id=?', id)
+        Shard.where("database_server_id IS NULL OR database_server_id=?", id)
       else
         Shard.where(database_server_id: id)
       end
@@ -187,7 +223,7 @@ module Switchman
       end
 
       id ||= begin
-        id_seq = Shard.connection.quote(Shard.connection.quote_table_name('switchman_shards_id_seq'))
+        id_seq = Shard.connection.quote(Shard.connection.quote_table_name("switchman_shards_id_seq"))
         next_id = Shard.connection.select_value("SELECT nextval(#{id_seq})")
         next_id.to_i
       end
@@ -204,29 +240,32 @@ module Switchman
                                 name: name,
                                 database_server_id: self.id)
           if create_statement
-            if ::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}")
+            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'
+              raise "This schema already exists; cannot overwrite"
             end
             Array(create_statement.call).each do |stmt|
               ::ActiveRecord::Base.connection.execute(stmt)
             end
           end
-          if config[:adapter] == 'postgresql'
-            old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor do
-            end
+          if config[:adapter] == "postgresql"
+            old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor {}
           end
           old_verbose = ::ActiveRecord::Migration.verbose
           ::ActiveRecord::Migration.verbose = false
 
           unless schema == false
             shard.activate do
-              reset_column_information
-
               ::ActiveRecord::Base.connection.transaction(requires_new: true) do
-                ::ActiveRecord::Base.connection.migration_context.migrate
+                if ::Rails.version < "7.1"
+                  ::ActiveRecord::Base.connection.migration_context.migrate
+                else
+                  ::ActiveRecord::MigrationContext.new(::ActiveRecord::Migrator.migrations_paths).migrate
+                end
               end
-              reset_column_information
+
               ::ActiveRecord::Base.descendants.reject do |m|
                 m <= UnshardedRecord || !m.table_exists?
               end.each(&:define_attribute_methods)
@@ -240,7 +279,6 @@ module Switchman
       rescue
         shard&.destroy
         shard&.drop_database rescue nil unless schema_already_existed
-        reset_column_information unless schema == false rescue nil
         raise
       ensure
         self.class.creating_new_shard = false
@@ -265,18 +303,18 @@ module Switchman
     end
 
     def primary_shard
-      unless instance_variable_defined?(:@primary_shard)
-        # if sharding isn't fully set up yet, we may not be able to query the shards table
-        @primary_shard = Shard.default if Shard.default.database_server == self
-        @primary_shard ||= shards.where(name: nil).first
-      end
-      @primary_shard
-    end
+      return nil unless primary_shard_id
 
-    private
+      Shard.lookup(primary_shard_id)
+    end
 
-    def reset_column_information
-      ::ActiveRecord::Base.descendants.reject { |m| m <= UnshardedRecord }.each(&:reset_column_information)
+    def primary_shard_id
+      unless instance_variable_defined?(:@primary_shard_id)
+        # if sharding isn't fully set up yet, we may not be able to query the shards table
+        @primary_shard_id = Shard.default.id if Shard.default.database_server == self
+        @primary_shard_id ||= shards.where(name: nil).first&.id
+      end
+      @primary_shard_id
     end
   end
 end

data/lib/switchman/default_shard.rb

@@ -3,9 +3,10 @@
 module Switchman
   class DefaultShard
     def id
-      'default'
+      "default"
     end
-    alias cache_key id
+    alias_method :cache_key, :id
+
     def activate(*_classes)
       yield
     end
@@ -57,8 +58,18 @@ module Switchman
       self
     end
 
+    def region; end
+
+    def in_region?(_region)
+      true
+    end
+
+    def in_current_region?
+      true
+    end
+
     def _dump(_depth)
-      ''
+      ""
     end
 
     def self._load(_str)

data/lib/switchman/engine.rb

@@ -5,24 +5,17 @@ module Switchman
     isolate_namespace Switchman
 
     # enable Rails 6.1 style connection handling
-    config.active_record.legacy_connection_handling = false
+    config.active_record.legacy_connection_handling = false if ::Rails.version < "7.1"
     config.active_record.writing_role = :primary
 
     ::GuardRail.singleton_class.prepend(GuardRail::ClassMethods)
 
-    # 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.active_record_patch',
-                before: 'active_record.initialize_database',
-                after: (::Rails.version < '7.0' ? :initialize_dependency_mechanism : :setup_once_autoloader) do
+    initializer "switchman.active_record_patch",
+                before: "active_record.initialize_database",
+                after: :setup_once_autoloader do
       ::ActiveSupport.on_load(:active_record) do
         # Switchman requires postgres, so just always load the pg adapter
-        require 'active_record/connection_adapters/postgresql_adapter'
+        require "active_record/connection_adapters/postgresql_adapter"
 
         self.default_shard = ::Rails.env.to_sym
         self.default_role = :primary
@@ -50,14 +43,21 @@ module Switchman
         ::ActiveRecord::Associations::CollectionProxy.include(ActiveRecord::Associations::CollectionProxy)
 
         ::ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecord::Associations::Preloader::Association)
-        ::ActiveRecord::Associations::Preloader::Association::LoaderRecords.prepend(ActiveRecord::Associations::Preloader::Association::LoaderRecords) unless ::Rails.version < '7.0'
+        ::ActiveRecord::Associations::Preloader::Association::LoaderRecords.prepend(
+          ActiveRecord::Associations::Preloader::Association::LoaderRecords
+        )
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::AbstractAdapter)
+        unless ::Rails.version < "7.1"
+          ::ActiveRecord::ConnectionAdapters::ConnectionHandler.prepend(ActiveRecord::ConnectionHandler)
+        end
         ::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)
+        ::ActiveRecord::DatabaseConfigurations::DatabaseConfig.prepend(
+          ActiveRecord::DatabaseConfigurations::DatabaseConfig
+        )
 
         ::ActiveRecord::LogSubscriber.prepend(ActiveRecord::LogSubscriber)
         ::ActiveRecord::Migration.prepend(ActiveRecord::Migration)
@@ -65,6 +65,11 @@ module Switchman
         ::ActiveRecord::MigrationContext.prepend(ActiveRecord::MigrationContext)
         ::ActiveRecord::Migrator.prepend(ActiveRecord::Migrator)
 
+        if ::Rails.version > "7.1.3"
+          ::ActiveRecord::PendingMigrationConnection.singleton_class
+                                                    .include(ActiveRecord::PendingMigrationConnection::ClassMethods)
+        end
+
         ::ActiveRecord::Reflection::AbstractReflection.include(ActiveRecord::Reflection::AbstractReflection)
         ::ActiveRecord::Reflection::AssociationReflection.prepend(ActiveRecord::Reflection::AssociationScopeCache)
         ::ActiveRecord::Reflection::ThroughReflection.prepend(ActiveRecord::Reflection::AssociationScopeCache)
@@ -77,8 +82,9 @@ module Switchman
         ::ActiveRecord::Relation.include(ActiveRecord::SpawnMethods)
         ::ActiveRecord::Relation.include(CallSuper)
 
-        ::ActiveRecord::PredicateBuilder::AssociationQueryValue.prepend(ActiveRecord::PredicateBuilder::AssociationQueryValue)
-        ::ActiveRecord::PredicateBuilder::PolymorphicArrayValue.prepend(ActiveRecord::PredicateBuilder::AssociationQueryValue)
+        ::ActiveRecord::PredicateBuilder::PolymorphicArrayValue.prepend(
+          ActiveRecord::PredicateBuilder::PolymorphicArrayValue
+        )
 
         ::ActiveRecord::Tasks::DatabaseTasks.singleton_class.prepend(ActiveRecord::Tasks::DatabaseTasks)
 
@@ -96,17 +102,18 @@ module Switchman
 
         ::ActiveRecord::ConnectionAdapters::TableDefinition.prepend(ActiveRecord::TableDefinition)
       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
+      # 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.error_patch', after: 'active_record.initialize_database' do
+    initializer "switchman.error_patch", after: "active_record.initialize_database" do
       ::ActiveSupport.on_load(:active_record) do
         ::StandardError.include(StandardError)
       end
     end
 
-    initializer 'switchman.initialize_cache', before: :initialize_cache, after: 'active_record.initialize_database' do
+    initializer "switchman.initialize_cache", before: :initialize_cache, after: "active_record.initialize_database" do
       ::ActiveSupport::Cache.singleton_class.prepend(ActiveSupport::Cache::ClassMethods)
 
       # if we haven't already setup our cache map out-of-band, set it up from
@@ -130,11 +137,11 @@ module Switchman
         Switchman.config[:cache_map][::Rails.env] = value
       end
 
-      middlewares = Switchman.config[:cache_map].values.map do |store|
+      middlewares = Switchman.config[:cache_map].values.filter_map do |store|
         store.middleware if store.respond_to?(:middleware)
-      end.compact.uniq
+      end.uniq
       middlewares.each do |middleware|
-        config.middleware.insert_before('Rack::Runtime', middleware)
+        config.middleware.insert_before("Rack::Runtime", middleware)
       end
 
       # prevent :initialize_cache from trying to (or needing to) set

data/lib/switchman/environment.rb

@@ -1,10 +1,10 @@
 # frozen_string_literal: true
 
-require 'etc'
+require "etc"
 
 module Switchman
   class Environment
-    def self.cpu_count(nproc_bin = 'nproc')
+    def self.cpu_count(nproc_bin = "nproc")
       return Etc.nprocessors if Etc.respond_to?(:nprocessors)
 
       `#{nproc_bin}`.to_i

data/lib/switchman/errors.rb

@@ -2,8 +2,21 @@
 
 module Switchman
   module Errors
+    class ManuallyCreatedShadowRecordError < RuntimeError
+      DEFAULT_MSG = "It looks like you're trying to manually create a shadow record. " \
+                    "Please use Switchman::ActiveRecord::Base#save_shadow_record instead."
+
+      def initialize(msg = DEFAULT_MSG)
+        super
+      end
+    end
+
     class NonExistentShardError < RuntimeError; end
 
     class ParallelShardExecError < RuntimeError; end
+
+    class ShadowRecordError < RuntimeError; end
+
+    class UnshardedTableError < RuntimeError; end
   end
 end

data/lib/switchman/guard_rail/relation.rb

@@ -14,7 +14,7 @@ module Switchman
 
       %w[update_all delete_all].each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
-          def #{method}(*args)
+          def #{method}(*args, **kwargs)
             db = Shard.current(connection_class_for_self).database_server
             db.unguard { super }
           end

data/lib/switchman/parallel.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'parallel'
+require "parallel"
 
 module Switchman
   module Parallel
@@ -50,19 +50,19 @@ module Switchman
       end
     end
 
-    class PrefixingIO
+    class TransformingIO
       delegate_missing_to :@original_io
 
-      def initialize(prefix, original_io)
-        @prefix = prefix
+      def initialize(transformer, original_io)
+        @transformer = transformer
         @original_io = original_io
       end
 
       def puts(*args)
-        args.flatten.each { |arg| @original_io.puts "#{@prefix}: #{arg}" }
+        args.flatten.each { |arg| @original_io.puts @transformer.call(arg) }
       end
     end
   end
 end
 
-::Parallel::UndumpableException.prepend(::Switchman::Parallel::UndumpableException)
+Parallel::UndumpableException.prepend(Switchman::Parallel::UndumpableException)

data/lib/switchman/r_spec_helper.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'switchman/test_helper'
+require "switchman/test_helper"
 
 module Switchman
   # including this module in your specs will give you several shards to
@@ -34,9 +34,9 @@ module Switchman
         groups = group.class.descendant_filtered_examples.map(&:example_group).uniq
         next unless groups.any? { |descendant_group| RSpecHelper.included_in?(descendant_group) }
 
-        puts 'Setting up sharding for all specs...'
+        puts "Setting up sharding for all specs..."
         Shard.delete_all
-        Switchman.cache.delete('default_shard')
+        Switchman.cache.delete("default_shard")
 
         @@shard1, @@shard2 = TestHelper.recreate_persistent_test_shards
         @@default_shard = Shard.default
@@ -48,7 +48,7 @@ module Switchman
             @@shard1 = @@shard1.create_new_shard
             @@shard2 = @@shard2.create_new_shard
           rescue => e
-            warn 'Sharding setup FAILED!:'
+            warn "Sharding setup FAILED!:"
             while e
               warn "\n#{e}\n"
               warn e.backtrace
@@ -66,9 +66,9 @@ module Switchman
         # we'll re-persist in the group's `before :all`; we don't want them to exist
         # in the db before then
         Shard.delete_all
-        Switchman.cache.delete('default_shard')
+        Switchman.cache.delete("default_shard")
         Shard.default(reload: true)
-        puts 'Done!'
+        puts "Done!"
 
         main_pid = Process.pid
         at_exit do
@@ -76,7 +76,7 @@ module Switchman
 
           # preserve rspec's exit status
           status = $!.is_a?(::SystemExit) ? $!.status : nil
-          puts 'Tearing down sharding for all specs'
+          puts "Tearing down sharding for all specs"
           @@shard1.database_server.destroy unless @@shard1.database_server == Shard.default.database_server
           unless @@keep_the_shards
             @@shard1.drop_database
@@ -95,7 +95,7 @@ module Switchman
         dup = @@default_shard.dup
         dup.id = @@default_shard.id
         dup.save!
-        Switchman.cache.delete('default_shard')
+        Switchman.cache.delete("default_shard")
         Shard.default(reload: true)
         dup = @@shard1.dup
         dup.id = @@shard1.id
@@ -107,7 +107,7 @@ module Switchman
       end
 
       klass.before do
-        raise 'Sharding did not set up correctly' if @@sharding_failed
+        raise "Sharding did not set up correctly" if @@sharding_failed
 
         Shard.clear_cache
         if use_transactional_tests
@@ -121,18 +121,19 @@ module Switchman
         next if @@sharding_failed
 
         # clean up after specs
-        DatabaseServer.all.each do |ds|
+        DatabaseServer.each do |ds|
           if ds.fake? && ds != @shard2.database_server
             ds.shards.delete_all unless use_transactional_tests
             ds.destroy
           end
+          ds.remove_instance_variable(:@primary_shard_id) if ds.instance_variable_defined?(:@primary_shard_id)
         end
       end
 
       klass.after(:all) do
         # Don't truncate because that can create some fun cross-connection lock contention
         Shard.delete_all
-        Switchman.cache.delete('default_shard')
+        Switchman.cache.delete("default_shard")
         Shard.default(reload: true)
       end
     end