switchman 3.5.9 → 3.5.16

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 64f3ba29604ff0a4f1db0b5c49000a636865f38c45ad6cb15eae252d98623d49
-  data.tar.gz: 68a22ead418800884673abd7a08ab84a38f548b065930e28258a25947ac4ab8e
+  metadata.gz: a27299a15c04348fc6a632a8c3c53b8c4fa541140b8befd07f63bc3544f01f50
+  data.tar.gz: d9a43fc82c874a87026c5d0b9467a45ba71df19e0c94166f22cc1aabe11b8e5f
 SHA512:
-  metadata.gz: 49bf3b77179a2d5a5788cfaec67aababa9bece27589bddadc387f2eafc2651b1cd6ac6397431d5a955fad05747acd8af2fe6e1c59597a447bf91209508771138
-  data.tar.gz: ed4ae505a24d605802c67793a69d8462905ed16f638490e8fa0bdca1bb8e55f2b49f154e9c7ebd85ac6d2ba92321e7de962a559a5489e7821cc094da944d207f
+  metadata.gz: 2b47b69ec2830ed31204d0fe122d555b9f6623683469d6b35329cc1bfe1a18b6e5ec769985d792f64c72a696a4a16f4119407a9323b2560f4be575f52d5c09b9
+  data.tar.gz: 19bd4da4fa0ba4e4ec03c57858481f6dddd799b98bcd9e67e3870bbfb60c7cb69d44ba74781788c084069dc5bc17632f74b3c793d1964a8ef76122f2f78c144e

data/lib/switchman/active_record/abstract_adapter.rb

@@ -27,8 +27,10 @@ module Switchman
         quote_table_name(name)
       end
 
-      def schema_migration
-        ::ActiveRecord::SchemaMigration
+      if ::Rails.version < "7.1"
+        def schema_migration
+          ::ActiveRecord::SchemaMigration
+        end
       end
 
       protected

data/lib/switchman/active_record/associations.rb

@@ -298,11 +298,56 @@ module Switchman
           record.has_attribute?(reflection.foreign_key) && record.send(reflection.foreign_key) != key
         end
 
-        def save_belongs_to_association(reflection)
-          # this seems counter-intuitive, but the autosave code will assign to attribute bypassing switchman,
-          # after reading the id attribute _without_ bypassing switchman. So we need Shard.current for the
-          # category of the associated record to match Shard.current for the category of self
-          shard.activate(connection_class_for_self_for_reflection(reflection)) { super }
+        if ::Rails.version < "7.1"
+          def save_belongs_to_association(reflection)
+            # this seems counter-intuitive, but the autosave code will assign to attribute bypassing switchman,
+            # after reading the id attribute _without_ bypassing switchman. So we need Shard.current for the
+            # category of the associated record to match Shard.current for the category of self
+            shard.activate(connection_class_for_self_for_reflection(reflection)) { super }
+          end
+        else
+          def save_belongs_to_association(reflection)
+            association = association_instance_get(reflection.name)
+            return unless association&.loaded? && !association.stale_target?
+
+            record = association.load_target
+            return unless record && !record.destroyed?
+
+            autosave = reflection.options[:autosave]
+
+            if autosave && record.marked_for_destruction?
+              foreign_key = Array(reflection.foreign_key)
+              foreign_key.each { |key| self[key] = nil }
+              record.destroy
+            elsif autosave != false
+              if record.new_record? || (autosave && record.changed_for_autosave?)
+                saved = record.save(validate: !autosave)
+              end
+
+              if association.updated?
+                primary_key = Array(compute_primary_key(reflection, record)).map(&:to_s)
+                foreign_key = Array(reflection.foreign_key)
+
+                primary_key_foreign_key_pairs = primary_key.zip(foreign_key)
+                primary_key_foreign_key_pairs.each do |pk, fk|
+                  # Notable change: add relative_id_for here
+                  association_id = if record.class.sharded_column?(pk)
+                                     Shard.relative_id_for(
+                                       record._read_attribute(pk),
+                                       record.shard,
+                                       shard
+                                     )
+                                   else
+                                     record._read_attribute(pk)
+                                   end
+                  self[fk] = association_id unless self[fk] == association_id
+                end
+                association.loaded!
+              end
+
+              saved if autosave
+            end
+          end
         end
       end
     end

data/lib/switchman/active_record/base.rb

@@ -57,7 +57,12 @@ module Switchman
         end
 
         def clear_query_caches_for_current_thread
-          ::ActiveRecord::Base.connection_handler.connection_pool_list.each do |pool|
+          pools = if ::Rails.version < "7.1"
+                    ::ActiveRecord::Base.connection_handler.connection_pool_list
+                  else
+                    ::ActiveRecord::Base.connection_handler.connection_pool_list(:all)
+                  end
+          pools.each do |pool|
             pool.connection(switch_shard: false).clear_query_cache if pool.active_connection?
           end
         end
@@ -158,7 +163,15 @@ module Switchman
                    end
 
         @loaded_from_shard ||= Shard.current(self.class.connection_class_for_self)
-        readonly! if shadow_record? && !Switchman.config[:writable_shadow_records]
+        if shadow_record? && !Switchman.config[:writable_shadow_records]
+          @readonly = true
+          @readonly_from_shadow ||= true
+        end
+        super
+      end
+
+      def readonly!
+        @readonly_from_shadow = false
         super
       end
 
@@ -169,6 +182,10 @@ module Switchman
         pkey > Shard::IDS_PER_SHARD
       end
 
+      def canonical?
+        !shadow_record?
+      end
+
       def save_shadow_record(new_attrs: attributes, target_shard: Shard.current)
         return if target_shard == shard
 
@@ -253,9 +270,9 @@ module Switchman
         result
       end
 
-      def transaction(**kwargs, &block)
+      def transaction(...)
         shard.activate(self.class.connection_class_for_self) do
-          self.class.transaction(**kwargs, &block)
+          self.class.transaction(...)
         end
       end
 

data/lib/switchman/active_record/connection_handler.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module ConnectionHandler
+      def resolve_pool_config(config, connection_name, role, shard)
+        ret = super
+        # Make *all* pool configs use the same schema reflection
+        ret.schema_reflection = ConnectionHandler.global_schema_reflection
+        ret
+      end
+
+      def self.global_schema_reflection
+        @global_schema_reflection ||= ::ActiveRecord::ConnectionAdapters::SchemaReflection.new(nil)
+      end
+    end
+  end
+end

data/lib/switchman/active_record/connection_pool.rb

@@ -3,22 +3,24 @@
 module Switchman
   module ActiveRecord
     module ConnectionPool
-      def get_schema_cache(connection)
-        self.schema_cache ||= SharedSchemaCache.get_schema_cache(connection)
-        self.schema_cache.connection = connection
+      if ::Rails.version < "7.1"
+        def get_schema_cache(connection)
+          self.schema_cache ||= SharedSchemaCache.get_schema_cache(connection)
+          self.schema_cache.connection = connection
 
-        self.schema_cache
-      end
+          self.schema_cache
+        end
 
-      # rubocop:disable Naming/AccessorMethodName override method
-      def set_schema_cache(cache)
-        schema_cache = get_schema_cache(cache.connection)
+        # rubocop:disable Naming/AccessorMethodName override method
+        def set_schema_cache(cache)
+          schema_cache = get_schema_cache(cache.connection)
 
-        cache.instance_variables.each do |x|
-          schema_cache.instance_variable_set(x, cache.instance_variable_get(x))
+          cache.instance_variables.each do |x|
+            schema_cache.instance_variable_set(x, cache.instance_variable_get(x))
+          end
         end
+        # rubocop:enable Naming/AccessorMethodName override method
       end
-      # rubocop:enable Naming/AccessorMethodName override method
 
       def default_schema
         connection unless @schemas

data/lib/switchman/active_record/database_configurations.rb

@@ -7,14 +7,26 @@ module Switchman
       # 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") }
+      if ::Rails.version < "7.1"
+        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
+      else
+        def configs_for(include_hidden: false, name: nil, **)
+          res = super
+          if name && !include_hidden
+            return nil unless name.end_with?("primary")
+          elsif !include_hidden
+            return res.select { |config| config.name.end_with?("primary") }
+          end
+          res
         end
-        res
       end
 
       private

data/lib/switchman/active_record/finder_methods.rb

@@ -42,26 +42,56 @@ module Switchman
         primary_shard.activate { super }
       end
 
-      def exists?(conditions = :none)
-        conditions = conditions.id if ::ActiveRecord::Base === conditions
-        return false unless conditions
+      if ::Rails.version < "7.1"
+        def exists?(conditions = :none)
+          conditions = conditions.id if ::ActiveRecord::Base === conditions
+          return false unless conditions
 
-        relation = apply_join_dependency(eager_loading: false)
-        return false if ::ActiveRecord::NullRelation === relation
+          relation = apply_join_dependency(eager_loading: false)
+          return false if ::ActiveRecord::NullRelation === relation
 
-        relation = relation.except(:select, :order).select("1 AS one").limit(1)
+          relation = relation.except(:select, :order).select("1 AS one").limit(1)
 
-        case conditions
-        when Array, Hash
-          relation = relation.where(conditions)
-        else
-          relation = relation.where(table[primary_key].eq(conditions)) if conditions != :none
+          case conditions
+          when Array, Hash
+            relation = relation.where(conditions)
+          else
+            relation = relation.where(table[primary_key].eq(conditions)) if conditions != :none
+          end
+
+          relation.activate do |shard_rel|
+            return true if connection.select_value(shard_rel.arel, "#{name} Exists")
+          end
+          false
         end
+      else
+        def exists?(conditions = :none)
+          return false if @none
+
+          if Base === conditions
+            raise ArgumentError, <<-TEXT.squish
+              You are passing an instance of ActiveRecord::Base to `exists?`.
+              Please pass the id of the object by calling `.id`.
+            TEXT
+          end
 
-        relation.activate do |shard_rel|
-          return true if connection.select_value(shard_rel.arel, "#{name} Exists")
+          return false if !conditions || limit_value == 0 # rubocop:disable Style/NumericPredicate
+
+          if eager_loading?
+            relation = apply_join_dependency(eager_loading: false)
+            return relation.exists?(conditions)
+          end
+
+          relation = construct_relation_for_exists(conditions)
+          return false if relation.where_clause.contradiction?
+
+          relation.activate do |shard_rel|
+            return true if skip_query_cache_if_necessary do
+                             connection.select_rows(shard_rel.arel, "#{name} Exists?").size == 1
+                           end
+          end
+          false
         end
-        false
       end
     end
   end

data/lib/switchman/active_record/log_subscriber.rb

@@ -5,8 +5,10 @@ module Switchman
     module LogSubscriber
       # sadly, have to completely replace this
       def sql(event)
-        self.class.runtime += event.duration
-        return unless logger.debug?
+        if ::Rails.version < "7.1"
+          self.class.runtime += event.duration
+          return unless logger.debug?
+        end
 
         payload = event.payload
 
@@ -27,7 +29,11 @@ module Switchman
         end
 
         name = colorize_payload_name(name, payload[:name])
-        sql  = color(sql, sql_color(sql), true)
+        sql  = if ::Rails.version < "7.1"
+                 color(sql, sql_color(sql), true)
+               else
+                 color(sql, sql_color(sql), bold: true)
+               end
 
         debug "  #{name}  #{sql}#{binds}#{shard}"
       end

data/lib/switchman/active_record/migration.rb

@@ -29,7 +29,11 @@ module Switchman
         db_name_hash = Zlib.crc32(Shard.current.name)
         shard_name_hash = ::ActiveRecord::Migrator::MIGRATOR_SALT * db_name_hash
         # Store in internalmetadata to allow other tools to be able to lock out migrations
-        ::ActiveRecord::InternalMetadata[:migrator_advisory_lock_id] = shard_name_hash
+        if ::Rails.version < "7.1"
+          ::ActiveRecord::InternalMetadata[:migrator_advisory_lock_id] = shard_name_hash
+        else
+          ::ActiveRecord::InternalMetadata.new(connection)[:migrator_advisory_lock_id] = shard_name_hash
+        end
         shard_name_hash
       end
 
@@ -48,17 +52,30 @@ module Switchman
     module MigrationContext
       def migrate(...)
         connection = ::ActiveRecord::Base.connection
-        connection_pool = ::ActiveRecord::Base.connection_pool
-        previous_schema_cache = connection_pool.get_schema_cache(connection)
-        temporary_schema_cache = ::ActiveRecord::ConnectionAdapters::SchemaCache.new(connection)
+        schema_cache_holder = ::ActiveRecord::Base.connection_pool
+        schema_cache_holder = schema_cache_holder.schema_reflection if ::Rails.version >= "7.1"
+        previous_schema_cache = if ::Rails.version < "7.1"
+                                  schema_cache_holder.get_schema_cache(connection)
+                                else
+                                  schema_cache_holder.instance_variable_get(:@cache)
+                                end
+
+        if ::Rails.version < "7.1"
+          temporary_schema_cache = ::ActiveRecord::ConnectionAdapters::SchemaCache.new(connection)
+
+          reset_column_information
+          schema_cache_holder.set_schema_cache(temporary_schema_cache)
+        else
+          schema_cache_holder.instance_variable_get(:@cache)
 
-        reset_column_information
-        connection_pool.set_schema_cache(temporary_schema_cache)
+          reset_column_information
+          schema_cache_holder.clear!
+        end
 
         begin
           super(...)
         ensure
-          connection_pool.set_schema_cache(previous_schema_cache)
+          schema_cache_holder.set_schema_cache(previous_schema_cache)
           reset_column_information
         end
       end

data/lib/switchman/active_record/persistence.rb

@@ -34,6 +34,10 @@ module Switchman
         # When a shadow record is reloaded the real record is returned. So
         # we need to ensure the loaded_from_shard is set correctly after a reload.
         @loaded_from_shard = @shard
+        if @readonly_from_shadow
+          @readonly_from_shadow = false
+          @readonly = false
+        end
         res
       end
 

data/lib/switchman/active_record/query_cache.rb

@@ -5,28 +5,57 @@ module Switchman
     module QueryCache
       private
 
-      def cache_sql(sql, name, binds)
-        # have to include the shard id in the cache key because of switching dbs on the same connection
-        sql = "#{shard.id}::#{sql}"
-        @lock.synchronize do
-          result =
-            if query_cache[sql].key?(binds)
-              args = {
-                sql: sql,
-                binds: binds,
-                name: name,
-                connection_id: object_id,
-                cached: true,
-                type_casted_binds: -> { type_casted_binds(binds) }
-              }
-              ::ActiveSupport::Notifications.instrument(
-                "sql.active_record",
-                args
-              )
-              query_cache[sql][binds]
+      if ::Rails.version < "7.1"
+        def cache_sql(sql, name, binds)
+          # have to include the shard id in the cache key because of switching dbs on the same connection
+          sql = "#{shard.id}::#{sql}"
+          @lock.synchronize do
+            result =
+              if query_cache[sql].key?(binds)
+                args = {
+                  sql: sql,
+                  binds: binds,
+                  name: name,
+                  connection_id: object_id,
+                  cached: true,
+                  type_casted_binds: -> { type_casted_binds(binds) }
+                }
+                ::ActiveSupport::Notifications.instrument(
+                  "sql.active_record",
+                  args
+                )
+                query_cache[sql][binds]
+              else
+                query_cache[sql][binds] = yield
+              end
+            result.dup
+          end
+        end
+      else
+        def cache_sql(sql, name, binds)
+          # have to include the shard id in the cache key because of switching dbs on the same connection
+          sql = "#{shard.id}::#{sql}"
+          key = binds.empty? ? sql : [sql, binds]
+          result = nil
+          hit = false
+
+          @lock.synchronize do
+            if (result = @query_cache.delete(key))
+              hit = true
+              @query_cache[key] = result
             else
-              query_cache[sql][binds] = yield
+              result = @query_cache[key] = yield
+              @query_cache.shift if @query_cache_max_size && @query_cache.size > @query_cache_max_size
             end
+          end
+
+          if hit
+            ::ActiveSupport::Notifications.instrument(
+              "sql.active_record",
+              cache_notification_info(sql, name, binds)
+            )
+          end
+
           result.dup
         end
       end

data/lib/switchman/active_record/query_methods.rb

@@ -175,7 +175,7 @@ module Switchman
       end
 
       def sharded_foreign_key?(relation, column)
-        models_for_table(relation.table_name).any? { |m| m.sharded_column?(column) }
+        models_for_table(relation.name).any? { |m| m.sharded_column?(column) }
       end
 
       def sharded_primary_key?(relation, column)
@@ -187,7 +187,7 @@ module Switchman
 
       def source_shard_for_foreign_key(relation, column)
         reflection = nil
-        models_for_table(relation.table_name).each do |model|
+        models_for_table(relation.name).each do |model|
           reflection = model.send(:reflection_for_integer_attribute, column)
           break if reflection
         end

data/lib/switchman/active_record/test_fixtures.rb

@@ -14,21 +14,28 @@ module Switchman
         # Code adapted from the code in rails proper
         @connection_subscriber =
           ::ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
-            spec_name = payload[:spec_name] if payload.key?(:spec_name)
+            spec_name = if ::Rails.version < "7.1"
+                          payload[:spec_name] if payload.key?(:spec_name)
+                        elsif payload.key?(:connection_name)
+                          payload[:connection_name]
+                        end
             shard = payload[:shard] if payload.key?(:shard)
-            setup_shared_connection_pool
 
             if spec_name && !FORBIDDEN_DB_ENVS.include?(shard)
               begin
                 connection = ::ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard: shard)
+                connection.connect! if ::Rails.version >= "7.1" # eagerly validate the connection
               rescue ::ActiveRecord::ConnectionNotEstablished, ::ActiveRecord::NoDatabaseError
                 connection = nil
               end
 
-              if connection && !@fixture_connections.include?(connection)
-                connection.begin_transaction joinable: false, _lazy: false
-                connection.pool.lock_thread = true if lock_threads
-                @fixture_connections << connection
+              if connection
+                setup_shared_connection_pool
+                unless @fixture_connections.include?(connection)
+                  connection.begin_transaction joinable: false, _lazy: false
+                  connection.pool.lock_thread = true if lock_threads
+                  @fixture_connections << connection
+                end
               end
             end
           end
@@ -37,7 +44,7 @@ module Switchman
       def enlist_fixture_connections
         setup_shared_connection_pool
 
-        ::ActiveRecord::Base.connection_handler.connection_pool_list.reject do |cp|
+        ::ActiveRecord::Base.connection_handler.connection_pool_list(:primary).reject do |cp|
           FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym)
         end.map(&:connection)
       end

data/lib/switchman/arel.rb

@@ -13,6 +13,20 @@ 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 << " "
@@ -24,27 +38,29 @@ module Switchman
           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/database_server.rb

@@ -10,6 +10,10 @@ module Switchman
       attr_accessor :creating_new_shard
       attr_reader :all_roles
 
+      include Enumerable
+
+      delegate :each, to: :all
+
       def all
         database_servers.values
       end
@@ -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)
@@ -140,6 +148,29 @@ 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)
@@ -228,7 +259,11 @@ module Switchman
           unless schema == false
             shard.activate do
               ::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
 
               ::ActiveRecord::Base.descendants.reject do |m|

data/lib/switchman/default_shard.rb

@@ -6,6 +6,7 @@ module Switchman
       "default"
     end
     alias_method :cache_key, :id
+
     def activate(*_classes)
       yield
     end
@@ -57,6 +58,16 @@ module Switchman
       self
     end
 
+    def region; end
+
+    def in_region?(_region)
+      true
+    end
+
+    def in_current_region?
+      true
+    end
+
     def _dump(_depth)
       ""
     end

data/lib/switchman/engine.rb

@@ -5,7 +5,7 @@ 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)
@@ -60,6 +60,9 @@ module Switchman
           )
         end
         ::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)

data/lib/switchman/r_spec_helper.rb

@@ -121,7 +121,7 @@ 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

data/lib/switchman/shard.rb

@@ -13,6 +13,47 @@ module Switchman
 
     scope :primary, -> { where(name: nil).order(:database_server_id, :id).distinct_on(:database_server_id) }
 
+    scope :in_region, (lambda do |region, include_regionless: true|
+      next in_current_region if region.nil?
+
+      dbs_by_region = DatabaseServer.group_by(&:region)
+      db_count_in_this_region = dbs_by_region[region]&.length.to_i
+      db_count_in_this_region += dbs_by_region[nil]&.length.to_i if include_regionless
+      non_existent_database_servers = Shard.send(:non_existent_database_servers)
+      db_count_in_other_regions = DatabaseServer.all.length -
+                                  db_count_in_this_region +
+                                  non_existent_database_servers.length
+
+      dbs_in_this_region = dbs_by_region[region]&.map(&:id) || []
+      dbs_in_this_region += dbs_by_region[nil]&.map(&:id) || [] if include_regionless
+
+      if db_count_in_this_region <= db_count_in_other_regions
+        if dbs_in_this_region.include?(Shard.default.database_server.id)
+          where("database_server_id IN (?) OR database_server_id IS NULL", dbs_in_this_region)
+        else
+          where(database_server_id: dbs_in_this_region)
+        end
+      elsif db_count_in_other_regions.zero?
+        all
+      else
+        dbs_not_in_this_region = DatabaseServer.map(&:id) - dbs_in_this_region + non_existent_database_servers
+        if dbs_in_this_region.include?(Shard.default.database_server.id)
+          where("database_server_id NOT IN (?) OR database_server_id IS NULL", dbs_not_in_this_region)
+        else
+          where.not(database_server_id: dbs_not_in_this_region)
+        end
+      end
+    end)
+
+    scope :in_current_region, (lambda do |include_regionless: true|
+      # sharding isn't set up? maybe we're in tests, or a somehow degraded environment
+      # either way there's only one shard, and we always want to see it
+      return [default] unless default.is_a?(Switchman::Shard)
+      return all if !Switchman.region || DatabaseServer.none?(&:region)
+
+      in_region(Switchman.region, include_regionless: include_regionless)
+    end)
+
     class << self
       def sharded_models
         @sharded_models ||= [::ActiveRecord::Base, UnshardedRecord].freeze
@@ -166,7 +207,11 @@ module Switchman
           # clear connections prior to forking (no more queries will be executed in the parent,
           # and we want them gone so that we don't accidentally use them post-fork doing something
           # silly like dealloc'ing prepared statements)
-          ::ActiveRecord::Base.clear_all_connections!
+          if ::Rails.version < "7.1"
+            ::ActiveRecord::Base.clear_all_connections!(nil)
+          else
+            ::ActiveRecord::Base.connection_handler.clear_all_connections!(:all)
+          end
 
           parent_process_name = sanitized_process_title
           ret = ::Parallel.map(scopes, in_processes: (scopes.length > 1) ? parallel : 0) do |server, subscope|
@@ -258,7 +303,7 @@ module Switchman
             else
               shard = partition_object.shard
             end
-          when Integer, /^\d+$/, /^(\d+)~(\d+)$/
+          when Integer, /\A\d+\Z/, /\A(\d+)~(\d+)\Z/
             local_id, shard = Shard.local_id_for(partition_object)
             local_id ||= partition_object
             object = local_id unless partition_proc
@@ -300,14 +345,14 @@ module Switchman
         case any_id
         when ::ActiveRecord::Base
           any_id.id
-        when /^(\d+)~(-?\d+)$/
+        when /\A(\d+)~(-?\d+)\Z/
           local_id = $2.to_i
           signed_id_operation(local_id) do |id|
             return nil if id > IDS_PER_SHARD
 
             ($1.to_i * IDS_PER_SHARD) + id
           end
-        when Integer, /^-?\d+$/
+        when Integer, /\A-?\d+\Z/
           any_id.to_i
         end
       end
@@ -391,7 +436,7 @@ module Switchman
       end
 
       def configure_connects_to
-        full_connects_to_hash = DatabaseServer.all.to_h { |db| [db.id.to_sym, db.connects_to_hash] }
+        full_connects_to_hash = DatabaseServer.to_h { |db| [db.id.to_sym, db.connects_to_hash] }
         sharded_models.each do |klass|
           connects_to_hash = full_connects_to_hash.deep_dup
           if klass == UnshardedRecord
@@ -410,7 +455,10 @@ module Switchman
             end
           end
 
+          # this resets the default shard on rails 7.1+, but we want to preserve it
+          shard_was = klass.default_shard
           klass.connects_to shards: connects_to_hash
+          klass.default_shard = shard_was
         end
       end
 
@@ -477,8 +525,17 @@ module Switchman
         argv[0] = File.basename(argv[0])
         argv.shelljoin
       end
+
+      # @return [Array<String>] the list of database servers that are in the
+      #   config, but don't have any shards on them
+      def non_existent_database_servers
+        @non_existent_database_servers ||=
+          Shard.distinct.pluck(:database_server_id).compact - DatabaseServer.all.map(&:id)
+      end
     end
 
+    delegate :region, :in_region?, :in_current_region?, to: :database_server
+
     def name
       unless instance_variable_defined?(:@name)
         # protect against re-entrancy
@@ -519,6 +576,10 @@ module Switchman
       id
     end
 
+    def original_id_value
+      id
+    end
+
     def activate(*classes, &block)
       shards = hashify_classes(classes)
       Shard.activate(shards, &block)

data/lib/switchman/standard_error.rb

@@ -13,8 +13,10 @@ module Switchman
         Thread.current[:switchman_error_handler] = true
 
         @active_shards ||= Shard.active_shards
-      rescue ThreadError # e.g. `require': can't be called from trap context (ThreadError)
-        # intentionally empty
+      rescue
+        # intentionally empty - don't allow calculating the active_shards to prevent
+        # creating the StandardError for any reason. this prevents various random issues
+        # when a StandardError is created within a finalizer
       ensure
         Thread.current[:switchman_error_handler] = nil
       end

data/lib/switchman/test_helper.rb

@@ -19,7 +19,7 @@ module Switchman
         end
 
         server1 = Shard.default.database_server
-        server2 = DatabaseServer.create(Shard.default.database_server.config.merge(server2: true))
+        server2 = DatabaseServer.create(Shard.default.database_server.config.merge(server2: true, schema_dump: false))
 
         if server1 == Shard.default.database_server && server1.config[:shard1] && server1.config[:shard2]
           # look for the shards in the db already

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = "3.5.9"
+  VERSION = "3.5.16"
 end

data/lib/switchman.rb

@@ -20,27 +20,31 @@ loader.setup
 module Switchman
   Deprecation = ::ActiveSupport::Deprecation.new("4.0", "Switchman")
 
-  def self.config
-    # TODO: load from yaml
-    @config ||= {}
-  end
+  class << self
+    attr_writer :cache
 
-  def self.cache
-    (@cache.respond_to?(:call) ? @cache.call : @cache) || ::Rails.cache
-  end
+    def config
+      # TODO: load from yaml
+      @config ||= {}
+    end
 
-  def self.cache=(cache)
-    @cache = cache
-  end
+    def cache
+      (@cache.respond_to?(:call) ? @cache.call : @cache) || ::Rails.cache
+    end
+
+    def region
+      config[:region]
+    end
 
-  def self.foreign_key_check(name, type, limit: nil)
-    return unless name.to_s.end_with?("_id") && type.to_s == "integer" && limit.to_i < 8
+    def foreign_key_check(name, type, limit: nil)
+      return unless name.to_s.end_with?("_id") && type.to_s == "integer" && limit.to_i < 8
 
-    puts <<~TEXT.squish
-      WARNING: All foreign keys need to be 8-byte integers.
-      #{name} looks like a foreign key.
-      If so, please add the option: `:limit => 8`
-    TEXT
+      puts <<~TEXT.squish
+        WARNING: All foreign keys need to be 8-byte integers.
+        #{name} looks like a foreign key.
+        If so, please add the option: `:limit => 8`
+      TEXT
+    end
   end
 
   class OrderOnMultiShardQuery < RuntimeError; end

data/lib/tasks/switchman.rake

@@ -35,7 +35,7 @@ module Switchman
 
         servers = servers.filter_map { |server| DatabaseServer.find(server) }
         if open
-          open_servers = DatabaseServer.all.select { |server| server.config[:open] }
+          open_servers = DatabaseServer.select { |server| server.config[:open] }
           servers.concat(open_servers)
           servers << DatabaseServer.find(nil) if open_servers.empty?
           servers.uniq!
@@ -43,6 +43,19 @@ module Switchman
         servers = DatabaseServer.all - servers if negative
       end
 
+      ENV["REGION"]&.split(",")&.each do |region|
+        method = :select!
+        if region[0] == "-"
+          method = :reject!
+          region = region[1..]
+        end
+        if region == "self"
+          servers.send(method, &:in_current_region?)
+        else
+          servers.send(method) { |server| server.in_region?(region) }
+        end
+      end
+
       servers = filter_database_servers_chain.call(servers)
 
       scope = base_scope.order(::Arel.sql("database_server_id IS NOT NULL, database_server_id, id"))

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.5.9
+  version: 3.5.16
 platform: ruby
 authors:
 - Cody Cutrer
@@ -10,7 +10,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2023-08-23 00:00:00.000000000 Z
+date: 2023-10-17 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -21,7 +21,7 @@ dependencies:
         version: 6.1.4
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.1'
+        version: '7.2'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
@@ -31,7 +31,7 @@ dependencies:
         version: 6.1.4
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.1'
+        version: '7.2'
 - !ruby/object:Gem::Dependency
   name: guardrail
   requirement: !ruby/object:Gem::Requirement
@@ -69,7 +69,7 @@ dependencies:
         version: '6.1'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.1'
+        version: '7.2'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
@@ -79,35 +79,21 @@ dependencies:
         version: '6.1'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.1'
+        version: '7.2'
 - !ruby/object:Gem::Dependency
-  name: appraisal
+  name: debug
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.3'
+        version: '1.8'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.3'
-- !ruby/object:Gem::Dependency
-  name: byebug
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
+        version: '1.8'
 - !ruby/object:Gem::Dependency
   name: pg
   requirement: !ruby/object:Gem::Requirement
@@ -122,20 +108,6 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '1.2'
-- !ruby/object:Gem::Dependency
-  name: pry
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
@@ -269,6 +241,7 @@ files:
 - lib/switchman/active_record/attribute_methods.rb
 - lib/switchman/active_record/base.rb
 - lib/switchman/active_record/calculations.rb
+- lib/switchman/active_record/connection_handler.rb
 - lib/switchman/active_record/connection_pool.rb
 - lib/switchman/active_record/database_configurations.rb
 - lib/switchman/active_record/database_configurations/database_config.rb