switchman 4.0.1 → 4.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bd3fc41cd455d5855db84fd49d132dca58c824dbe50e49ea4ab7c4cd0ee24d00
-  data.tar.gz: df1df2b43f521a5b5e424165ac167d2de351c4800db6fe50b5d83a78e8e17e16
+  metadata.gz: e03940e1a5f3d0e7b933d99b3cbd68e3998b55f14ac0e7eec91ef81a193445e1
+  data.tar.gz: 923793689418332ab6ce2ffef2981567bd74b6358916794d958172638b0ec160
 SHA512:
-  metadata.gz: 5b721c57764dc47c2da10b027b63fdeed4795c349b2dec36b0d440e98ff0294355a4bbb03d2176f30883de0b842e2f4a47f27ea2fa2f56405594f53d4a278f6c
-  data.tar.gz: 2729dfb39b1e3d8e3df84bcf1bd321ac39de73e552d8a6c0f15f48f6b9fdb47611a076f8e11fbcbfeb4fd8097c2e6a49a2eab7d76942af92ed769c03b1b9f7d0
+  metadata.gz: 6950ac12cf6532f98e3674ded8890e7bb2f9ef50f33657220096de50508715ab047250c163546dc7396c61d15197ee2b517ae342f84cbe007ae22a217ed7751f
+  data.tar.gz: c45df067998ec5868bfe2c6a0f52608ffd9de9a19e2ce55ba55fc9986e077833df9964c5783c0f394a9683cb21a610e7f2a781d88e690fbc2bfeef5e8e352826

data/lib/switchman/active_record/abstract_adapter.rb

@@ -5,7 +5,7 @@ module Switchman
     module AbstractAdapter
       module ForeignKeyCheck
         def add_column(table, name, type, limit: nil, **)
-          Switchman.foreign_key_check(name, type, limit: limit)
+          Switchman.foreign_key_check(name, type, limit:)
           super
         end
       end

data/lib/switchman/active_record/associations.rb

@@ -174,7 +174,7 @@ module Switchman
           end
 
           # Disabling to keep closer to rails original
-          # rubocop:disable Naming/AccessorMethodName, Style/GuardClause
+          # rubocop:disable Naming/AccessorMethodName
           # significant changes:
           #  * globalize the key to lookup
           def set_inverse(record)
@@ -191,7 +191,7 @@ module Switchman
               association.set_inverse_instance(record)
             end
           end
-          # rubocop:enable Naming/AccessorMethodName, Style/GuardClause
+          # rubocop:enable Naming/AccessorMethodName
 
           # significant changes:
           #  * partition_by_shard the records_for call

data/lib/switchman/active_record/attribute_methods.rb

@@ -59,9 +59,9 @@ module Switchman
           if ::Rails.version < "7.1.4"
             # https://github.com/rails/rails/commit/a2a12fc2e3f4e6d06f81d4c74c88f8e6b3369ee6#diff-5b59ece6d9396b596f06271cec0ea726e3360911383511c49b1a66f454bfc2b6L30
             # These arguments were effectively swapped in Rails 7.1.4, so previous versions need them reversed
-            owner.define_cached_method(as, namespace: namespace, as: name, &block)
+            owner.define_cached_method(as, namespace:, as: name, &block)
           else
-            owner.define_cached_method(name, namespace: namespace, as: as, &block)
+            owner.define_cached_method(name, namespace:, as:, &block)
           end
         end
 
@@ -136,14 +136,14 @@ module Switchman
             safe_class_name = class_name.unpack1("h*")
             define_cached_method(owner,
                                  "sharded_#{safe_class_name}_#{attr_name}",
-                                 as: as,
+                                 as:,
                                  namespace: :switchman) do |batch|
               batch << build_sharded_getter("sharded_#{safe_class_name}_#{attr_name}",
                                             "original_#{as}",
                                             class_name)
             end
           else
-            define_cached_method(owner, "plain_#{attr_name}", as: as, namespace: :switchman) do |batch|
+            define_cached_method(owner, "plain_#{attr_name}", as:, namespace: :switchman) do |batch|
               batch << <<-RUBY
                 def plain_#{attr_name}
                   _read_attribute("#{attr_name}") { |n| missing_attribute(n, caller) }

data/lib/switchman/active_record/base.rb

@@ -6,9 +6,9 @@ module Switchman
       module ClassMethods
         delegate :shard, to: :all
 
-        def find_ids_in_ranges(opts = {}, &block)
+        def find_ids_in_ranges(opts = {}, &)
           opts.reverse_merge!(loose: true)
-          all.find_ids_in_ranges(opts, &block)
+          all.find_ids_in_ranges(opts, &)
         end
 
         def sharded_model
@@ -34,56 +34,6 @@ module Switchman
           end
         end
 
-        # NOTE: `returning` values are _not_ transposed back to the current shard
-        %w[insert_all upsert_all].each do |method|
-          class_eval <<-RUBY, __FILE__, __LINE__ + 1
-            def #{method}(attributes, returning: nil, **)
-              scope = self != ::ActiveRecord::Base && current_scope
-              if (target_shard = scope&.primary_shard) == (current_shard = Shard.current(connection_class_for_self))
-                scope = nil
-              end
-              if scope
-                dupped = false
-                attributes.each_with_index do |hash, i|
-                  if dupped || hash.any? { |k, v| sharded_column?(k) }
-                    unless dupped
-                      attributes = attributes.dup
-                      dupped = true
-                    end
-                    attributes[i] = hash.to_h do |k, v|
-                      if sharded_column?(k)
-                        [k, Shard.relative_id_for(v, current_shard, target_shard)]
-                      else
-                        [k, v]
-                      end
-                    end
-                  end
-                end
-              end
-
-              if scope
-                scope.activate do
-                  db = Shard.current(connection_class_for_self).database_server
-                  result = db.unguard { super }
-                  if result&.columns&.any? { |c| sharded_column?(c) }
-                    transposed_rows = result.rows.map do |row|
-                      row.map.with_index do |value, i|
-                        sharded_column?(result.columns[i]) ? Shard.relative_id_for(value, target_shard, current_shard) : value
-                      end
-                    end
-                    result = ::ActiveRecord::Result.new(result.columns, transposed_rows, result.column_types)
-                  end
-
-                  result
-                end
-              else
-                db = Shard.current(connection_class_for_self).database_server
-                db.unguard { super }
-              end
-            end
-          RUBY
-        end
-
         def reset_column_information
           @sharded_column_values = {}
           super
@@ -109,7 +59,11 @@ module Switchman
                     ::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?
+            if ::Rails.version < "7.2"
+              pool.connection(switch_shard: false).clear_query_cache if pool.active_connection?
+            elsif pool.active_connection?
+              pool.lease_connection(switch_shard: false).clear_query_cache
+            end
           end
         end
 
@@ -132,7 +86,7 @@ module Switchman
                               :"#{current_shard}/#{current_role}"
                             end
 
-          super(config_or_env)
+          super
         end
 
         def connected_to_stack
@@ -189,6 +143,7 @@ module Switchman
 
       def self.prepended(klass)
         klass.singleton_class.prepend(ClassMethods)
+        klass.singleton_class.prepend(Switchman::ActiveRecord::Relation::InsertUpsertAll) if ::Rails.version < "7.2"
         klass.scope :non_shadow, lambda { |key = primary_key|
                                    where(key => (QueryMethods::NonTransposingValue.new(0)..
                                                  QueryMethods::NonTransposingValue.new(Shard::IDS_PER_SHARD)))
@@ -241,7 +196,7 @@ module Switchman
                                end
         end
         target_shard.activate do
-          self.class.upsert(shadow_attrs, unique_by: self.class.primary_key)
+          self.class.upsert_all([shadow_attrs], unique_by: self.class.primary_key)
         end
       end
 

data/lib/switchman/active_record/calculations.rb

@@ -159,7 +159,7 @@ module Switchman
       end
 
       def grouped_calculation_options(operation, column_name, distinct)
-        opts = { operation: operation, column_name: column_name, distinct: distinct }
+        opts = { operation:, column_name:, distinct: }
 
         # Rails 7.0.5
         if defined?(::ActiveRecord::Calculations::ColumnAliasTracker)
@@ -188,11 +188,11 @@ module Switchman
         group_columns = group_aliases.zip(group_fields).map do |aliaz, field|
           [aliaz, type_for(field), column_name_for(field)]
         end
-        opts.merge!(association: association,
-                    associated: associated,
-                    group_aliases: group_aliases,
-                    group_columns: group_columns,
-                    group_fields: group_fields)
+        opts.merge!(association:,
+                    associated:,
+                    group_aliases:,
+                    group_columns:,
+                    group_fields:)
 
         opts
       end

data/lib/switchman/active_record/connection_pool.rb

@@ -11,7 +11,7 @@ module Switchman
           self.schema_cache
         end
 
-        # rubocop:disable Naming/AccessorMethodName override method
+        # rubocop:disable Naming/AccessorMethodName -- override method
         def set_schema_cache(cache)
           schema_cache = get_schema_cache(cache.connection)
 
@@ -19,13 +19,14 @@ module Switchman
             schema_cache.instance_variable_set(x, cache.instance_variable_get(x))
           end
         end
-        # rubocop:enable Naming/AccessorMethodName override method
+        # rubocop:enable Naming/AccessorMethodName -- override method
       end
 
       def default_schema
-        connection unless @schemas
+        connection_method = (::Rails.version < "7.2") ? :connection : :lease_connection
+        send(connection_method) unless @schemas
         # default shard will not switch databases immediately, so it won't be set yet
-        @schemas ||= connection.current_schemas
+        @schemas ||= send(connection_method).current_schemas
         @schemas.first
       end
 
@@ -43,8 +44,36 @@ module Switchman
         conn
       end
 
+      unless ::Rails.version < "7.2"
+        def active_connection(switch_shard: true)
+          conn = super()
+          return nil if conn.nil?
+          raise Errors::NonExistentShardError if current_shard.new_record?
+
+          switch_database(conn) if conn.shard != current_shard && switch_shard
+          conn
+        end
+
+        def lease_connection(switch_shard: true)
+          conn = super()
+          raise Errors::NonExistentShardError if current_shard.new_record?
+
+          switch_database(conn) if conn.shard != current_shard && switch_shard
+          conn
+        end
+
+        def with_connection(switch_shard: true, **kwargs)
+          super(**kwargs) do |conn|
+            raise Errors::NonExistentShardError if current_shard.new_record?
+
+            switch_database(conn) if conn.shard != current_shard && switch_shard
+            yield conn
+          end
+        end
+      end
+
       def release_connection(with_id = Thread.current)
-        super(with_id)
+        super
 
         flush
       end

data/lib/switchman/active_record/finder_methods.rb

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module FinderMethods
       def find_one(id)
-        return super(id) unless klass.integral_id?
+        return super unless klass.integral_id?
 
         if shard_source_value != :implicit
           current_shard = Shard.current(klass.connection_class_for_self)
@@ -28,14 +28,14 @@ module Switchman
         if shard
           shard.activate { super(local_id) }
         else
-          super(id)
+          super
         end
       end
 
       def find_some_ordered(ids)
         current_shard = Shard.current(klass.connection_class_for_self)
         ids = ids.map { |id| Shard.relative_id_for(id, current_shard, current_shard) }
-        super(ids)
+        super
       end
 
       def find_or_instantiator_by_attributes(match, attributes, *args)

data/lib/switchman/active_record/migration.rb

@@ -31,8 +31,10 @@ module Switchman
         # Store in internalmetadata to allow other tools to be able to lock out migrations
         if ::Rails.version < "7.1"
           ::ActiveRecord::InternalMetadata[:migrator_advisory_lock_id] = shard_name_hash
-        else
+        elsif ::Rails.version < "7.2"
           ::ActiveRecord::InternalMetadata.new(connection)[:migrator_advisory_lock_id] = shard_name_hash
+        else
+          ::ActiveRecord::InternalMetadata.new(connection.pool)[:migrator_advisory_lock_id] = shard_name_hash
         end
         shard_name_hash
       end
@@ -73,9 +75,13 @@ module Switchman
         end
 
         begin
-          super(...)
+          super
         ensure
-          schema_cache_holder.set_schema_cache(previous_schema_cache)
+          if ::Rails.version < "7.2"
+            schema_cache_holder.set_schema_cache(previous_schema_cache)
+          else
+            schema_cache_holder.instance_variable_set(:@cache, previous_schema_cache)
+          end
           reset_column_information
         end
       end

data/lib/switchman/active_record/persistence.rb

@@ -24,7 +24,7 @@ module Switchman
         super
       end
 
-      def create_or_update(**, &block)
+      def create_or_update(**, &)
         writable_shadow_record_warning
         super
       end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -73,23 +73,37 @@ module Switchman
         end
       end
 
+      module ClassMethods
+        def quote_local_table_name(name)
+          # postgres quotes tables and columns the same; just pass through
+          # (differs from quote_table_name_with_shard below by no logic to
+          # explicitly qualify the table)
+          quote_column_name(name)
+        end
+
+        def quote_table_name(name, shard: nil)
+          # This looks kind of weird at first glance, but older Rails versions do not actually import
+          # these methods as class methods.
+          shard = self.shard if shard.nil? && ::Rails.version < "7.2" && !@use_local_table_name
+
+          return quote_local_table_name(name) unless shard
+
+          name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
+          name.instance_variable_set(:@schema, shard.name) unless name.schema
+          name.quoted
+        end
+      end
+
       def quote_local_table_name(name)
-        # postgres quotes tables and columns the same; just pass through
-        # (differs from quote_table_name_with_shard below by no logic to
-        # explicitly qualify the table)
-        quote_column_name(name)
+        self.class.quote_local_table_name(name)
       end
 
       def quote_table_name(name)
-        return quote_local_table_name(name) if @use_local_table_name
-
-        name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
-        name.instance_variable_set(:@schema, shard.name) unless name.schema
-        name.quoted
+        self.class.quote_table_name(name, shard: @use_local_table_name ? nil : shard)
       end
 
-      def with_global_table_name(&block)
-        with_local_table_name(false, &block)
+      def with_global_table_name(&)
+        with_local_table_name(false, &)
       end
 
       def with_local_table_name(enable = true) # rubocop:disable Style/OptionalBooleanParameter
@@ -129,7 +143,7 @@ module Switchman
       end
 
       def columns(*)
-        with_local_table_name(false) { super }
+        with_global_table_name { super }
       end
     end
   end

data/lib/switchman/active_record/query_cache.rb

@@ -13,9 +13,9 @@ module Switchman
             result =
               if query_cache[sql].key?(binds)
                 args = {
-                  sql: sql,
-                  binds: binds,
-                  name: name,
+                  sql:,
+                  binds:,
+                  name:,
                   connection_id: object_id,
                   cached: true,
                   type_casted_binds: -> { type_casted_binds(binds) }
@@ -40,12 +40,20 @@ module Switchman
           hit = false
 
           @lock.synchronize do
-            if (result = @query_cache.delete(key))
-              hit = true
-              @query_cache[key] = result
+            if ::Rails.version < "7.2"
+              if (result = @query_cache.delete(key))
+                hit = true
+                @query_cache[key] = result
+              else
+                result = @query_cache[key] = yield
+                @query_cache.shift if @query_cache_max_size && @query_cache.size > @query_cache_max_size
+              end
             else
-              result = @query_cache[key] = yield
-              @query_cache.shift if @query_cache_max_size && @query_cache.size > @query_cache_max_size
+              hit = true
+              result = @query_cache.compute_if_absent(key) do
+                hit = false
+                yield
+              end
             end
           end
 

data/lib/switchman/active_record/query_methods.rb

@@ -34,13 +34,29 @@ module Switchman
       end
 
       def shard_value=(value)
-        raise ::ActiveRecord::ImmutableRelation if @loaded
+        if @loaded
+          error_class = if ::Rails.version < "7.2"
+                          ::ActiveRecord::ImmutableRelation
+                        else
+                          ::ActiveRecord::UnmodifiableRelation
+                        end
+
+          raise error_class
+        end
 
         @values[:shard] = value
       end
 
       def shard_source_value=(value)
-        raise ::ActiveRecord::ImmutableRelation if @loaded
+        if @loaded
+          error_class = if ::Rails.version < "7.2"
+                          ::ActiveRecord::ImmutableRelation
+                        else
+                          ::ActiveRecord::UnmodifiableRelation
+                        end
+
+          raise error_class
+        end
 
         @values[:shard_source] = value
       end
@@ -107,6 +123,10 @@ module Switchman
 
       protected
 
+      def arel_columns(columns)
+        connection.with_local_table_name { super }
+      end
+
       def remove_nonlocal_primary_keys!
         each_transposable_predicate_value do |value, predicate, _relation, _column, type|
           next value unless
@@ -243,31 +263,49 @@ module Switchman
         end
       end
 
-      def arel_columns(columns)
-        connection.with_local_table_name { super }
-      end
-
       def arel_column(columns)
         connection.with_local_table_name { super }
       end
 
       def table_name_matches?(from)
-        connection.with_global_table_name { super }
+        if ::Rails.version < "7.2"
+          connection.with_global_table_name { super }
+        else
+          connection.with_global_table_name do
+            table_name = Regexp.escape(table.name)
+            # INST: adapter_class -> connection
+            quoted_table_name = Regexp.escape(connection.quote_table_name(table.name))
+            /(?:\A|(?<!FROM)\s)(?:\b#{table_name}\b|#{quoted_table_name})(?!\.)/i.match?(from.to_s)
+          end
+        end
+      end
+
+      unless ::Rails.version < "7.2"
+        def order_column(field)
+          arel_column(field) do |attr_name|
+            if attr_name == "count" && !group_values.empty?
+              table[attr_name]
+            else
+              # INST: adapter_class -> connection
+              ::Arel.sql(connection.quote_table_name(attr_name), retryable: true)
+            end
+          end
+        end
       end
 
-      def each_predicate(predicates = nil, &block)
-        return predicates.map(&block) if predicates
+      def each_predicate(predicates = nil, &)
+        return predicates.map(&) if predicates
 
-        each_predicate_cb(:having_clause, :having_clause=, &block)
-        each_predicate_cb(:where_clause, :where_clause=, &block)
+        each_predicate_cb(:having_clause, :having_clause=, &)
+        each_predicate_cb(:where_clause, :where_clause=, &)
       end
 
-      def each_predicate_cb(clause_getter, clause_setter, &block)
+      def each_predicate_cb(clause_getter, clause_setter, &)
         old_clause = send(clause_getter)
         old_predicates = old_clause.send(:predicates)
         return if old_predicates.empty?
 
-        new_predicates = old_predicates.map(&block)
+        new_predicates = old_predicates.map(&)
         return if new_predicates == old_predicates
 
         new_clause = old_clause.dup
@@ -289,7 +327,10 @@ module Switchman
 
             next predicate if new_left == old_left && new_right == old_right
 
-            next predicate.class.new predicate.expr.class.new(new_left, new_right)
+            next predicate.class.new predicate.expr.class.new(new_left, new_right) if ::Rails.version < "7.2"
+
+            next predicate.class.new predicate.expr.class.new([new_left, new_right])
+
           when ::Arel::Nodes::SelectStatement
             new_cores = predicate.cores.map do |core|
               next core unless core.is_a?(::Arel::Nodes::SelectCore) # just in case something weird is going on
@@ -344,33 +385,33 @@ module Switchman
         end
       end
 
-      def each_transposable_predicate_value_cb(node, original_block, &block)
+      def each_transposable_predicate_value_cb(node, original_block, &)
         case node
         when Array
-          node.filter_map { |val| each_transposable_predicate_value_cb(val, original_block, &block).presence }
+          node.filter_map { |val| each_transposable_predicate_value_cb(val, original_block, &).presence }
         when ::ActiveModel::Attribute
           old_value = node.value_before_type_cast
-          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (old_value == new_value) ? node : node.class.new(node.name, new_value, node.type)
         when ::Arel::Nodes::And
           old_value = node.children
-          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (old_value == new_value) ? node : node.class.new(new_value)
         when ::Arel::Nodes::BindParam
           old_value = node.value
-          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (old_value == new_value) ? node : node.class.new(new_value)
         when ::Arel::Nodes::Casted
           old_value = node.value
-          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (old_value == new_value) ? node : node.class.new(new_value, node.attribute)
         when ::Arel::Nodes::HomogeneousIn
           old_value = node.values
-          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           # switch to a regular In, so that Relation::WhereClause#contradiction? knows about it
           if new_value.empty?
@@ -381,7 +422,7 @@ module Switchman
           end
         when ::Arel::Nodes::Binary
           old_value = node.right
-          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (old_value == new_value) ? node : node.class.new(node.left, new_value)
         when ::Arel::Nodes::SelectStatement

data/lib/switchman/active_record/reflection.rb

@@ -28,11 +28,18 @@ module Switchman
         # this technically belongs on AssociationReflection, but we put it on
         # ThroughReflection as well, instead of delegating to its internal
         # HasManyAssociation, losing its proper `klass`
-        def association_scope_cache(klass, owner, &block)
+        def association_scope_cache(klass, owner, &)
           key = self
           key = [key, owner._read_attribute(@foreign_type)] if polymorphic?
           key = [key, shard(owner).id].flatten
-          klass.cached_find_by_statement(key, &block)
+
+          if ::Rails.version < "7.2"
+            klass.cached_find_by_statement(key, &)
+          else
+            klass.with_connection do |connection|
+              klass.cached_find_by_statement(connection, key, &)
+            end
+          end
         end
       end
 

data/lib/switchman/active_record/relation.rb

@@ -28,15 +28,15 @@ module Switchman
         relation
       end
 
-      def new(*, &block)
+      def new(*, &)
         primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
-      def create(*, &block)
+      def create(*, &)
         primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
-      def create!(*, &block)
+      def create!(*, &)
         primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
@@ -73,6 +73,59 @@ module Switchman
         RUBY
       end
 
+      # https://github.com/rails/rails/commit/ed2c15b52450ff927a05629f031376f25b670335
+      # once the minimum version is Rails 7.2, we can drop this separate module
+      module InsertUpsertAll
+        %w[insert_all upsert_all].each do |method|
+          class_eval <<-RUBY, __FILE__, __LINE__ + 1
+            def #{method}(attributes, returning: nil, **)
+              scope = self != ::ActiveRecord::Base && current_scope
+              if (target_shard = scope&.primary_shard) == (current_shard = Shard.current(connection_class_for_self))
+                scope = nil
+              end
+              if scope
+                dupped = false
+                attributes.each_with_index do |hash, i|
+                  if dupped || hash.any? { |k, v| sharded_column?(k) }
+                    unless dupped
+                      attributes = attributes.dup
+                      dupped = true
+                    end
+                    attributes[i] = hash.to_h do |k, v|
+                      if sharded_column?(k)
+                        [k, Shard.relative_id_for(v, current_shard, target_shard)]
+                      else
+                        [k, v]
+                      end
+                    end
+                  end
+                end
+              end
+
+              if scope
+                scope.activate do
+                  db = Shard.current(connection_class_for_self).database_server
+                  result = db.unguard { super }
+                  if result&.columns&.any? { |c| sharded_column?(c) }
+                    transposed_rows = result.rows.map do |row|
+                      row.map.with_index do |value, i|
+                        sharded_column?(result.columns[i]) ? Shard.relative_id_for(value, target_shard, current_shard) : value
+                      end
+                    end
+                    result = ::ActiveRecord::Result.new(result.columns, transposed_rows, result.column_types)
+                  end
+
+                  result
+                end
+              else
+                db = Shard.current(connection_class_for_self).database_server
+                db.unguard { super }
+              end
+            end
+          RUBY
+        end
+      end
+
       def find_ids_in_ranges(options = {})
         is_integer = columns_hash[primary_key.to_s].type == :integer
         loose_mode = options[:loose] && is_integer

data/lib/switchman/active_record/spawn_methods.rb

@@ -62,16 +62,12 @@ module Switchman
         if primary_shard != final_primary_shard && rhs.primary_shard != final_primary_shard
           shard!(final_primary_shard)
           rhs = rhs.shard(final_primary_shard)
-          super(rhs)
         elsif primary_shard != final_primary_shard
           shard!(final_primary_shard)
-          super(rhs)
         elsif rhs.primary_shard != final_primary_shard
           rhs = rhs.shard(final_primary_shard)
-          super(rhs)
-        else
-          super
         end
+        super
 
         self.shard_value = final_shard_value
         self.shard_source_value = final_shard_source_value

data/lib/switchman/active_record/table_definition.rb

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

data/lib/switchman/active_record/test_fixtures.rb

@@ -4,49 +4,89 @@ module Switchman
   module ActiveRecord
     module TestFixtures
       FORBIDDEN_DB_ENVS = %i[development production].freeze
-      def setup_fixtures(config = ::ActiveRecord::Base)
-        super
-
-        return unless run_in_transaction?
-
-        # Replace the one that activerecord natively uses with a switchman-optimized one
-        ::ActiveSupport::Notifications.unsubscribe(@connection_subscriber)
-        # Code adapted from the code in rails proper
-        @connection_subscriber =
-          ::ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
-            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)
 
-            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 ::Rails.version < "7.2"
+        def setup_fixtures(config = ::ActiveRecord::Base)
+          super
+          return unless run_in_transaction?
 
-              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
+          # Replace the one that activerecord natively uses with a switchman-optimized one
+          ::ActiveSupport::Notifications.unsubscribe(@connection_subscriber)
+          # Code adapted from the code in rails proper
+          @connection_subscriber =
+            ::ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
+              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)
+
+              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
+                  connection = nil
+                end
+
+                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
+        end
+
+        def enlist_fixture_connections
+          setup_shared_connection_pool
+
+          ::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
+      else
+        def setup_transactional_fixtures
+          setup_shared_connection_pool
+
+          # Begin transactions for connections already established
+          # INST: :writing -> :primary
+          @fixture_connection_pools = ::ActiveRecord::Base.connection_handler.connection_pool_list(:primary)
+          # INST: filter by FORBIDDEN_DB_ENVS
+          @fixture_connection_pools = @fixture_connection_pools.reject do |cp|
+            FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym)
           end
-      end
 
-      def enlist_fixture_connections
-        setup_shared_connection_pool
+          @fixture_connection_pools.each do |pool|
+            pool.pin_connection!(lock_threads)
+            pool.lease_connection
+          end
+
+          # When connections are established in the future, begin a transaction too
+          @connection_subscriber = ::ActiveSupport::Notifications
+                                   .subscribe("!connection.active_record") do |_, _, _, _, payload|
+            connection_name = payload[:connection_name] if payload.key?(:connection_name)
+            shard = payload[:shard] if payload.key?(:shard)
+
+            # INST: filter by FORBIDDEN_DB_ENVS
+            if connection_name && !FORBIDDEN_DB_ENVS.include?(shard)
+              pool = ::ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_name, shard: shard)
+              if pool
+                setup_shared_connection_pool
 
-        ::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)
+                unless @fixture_connection_pools.include?(pool)
+                  pool.pin_connection!(lock_threads)
+                  pool.lease_connection
+                  @fixture_connection_pools << pool
+                end
+              end
+            end
+          end
+        end
       end
     end
   end

data/lib/switchman/database_server.rb

@@ -121,7 +121,7 @@ module Switchman
       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,
+                                                          role:,
                                                           shard: id.to_sym)
         end
       end
@@ -218,7 +218,7 @@ module Switchman
       if config_create_statement
         create_commands = Array(config_create_statement).dup
         create_statement = lambda {
-          create_commands.map { |statement| format(statement, name: name, password: password) }
+          create_commands.map { |statement| format(statement, name:, password:) }
         }
       end
 
@@ -236,8 +236,8 @@ module Switchman
         self.class.creating_new_shard = true
         DatabaseServer.send(:reference_role, :deploy)
         ::ActiveRecord::Base.connected_to(shard: self.id.to_sym, role: :deploy) do
-          shard = Shard.create!(id: id,
-                                name: name,
+          shard = Shard.create!(id:,
+                                name:,
                                 database_server_id: self.id)
           if create_statement
             if ::ActiveRecord::Base.connection.select_value(

data/lib/switchman/engine.rb

@@ -53,6 +53,14 @@ module Switchman
         ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ActiveRecord::ConnectionPool)
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::QueryCache)
         ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter)
+        # https://github.com/rails/rails/commit/0016280f4fde55d96738887093dc333aae0d107b
+        if ::Rails.version < "7.2"
+          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter::ClassMethods)
+        else
+          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.singleton_class.prepend(
+            ActiveRecord::PostgreSQLAdapter::ClassMethods
+          )
+        end
 
         ::ActiveRecord::DatabaseConfigurations.prepend(ActiveRecord::DatabaseConfigurations)
         ::ActiveRecord::DatabaseConfigurations::DatabaseConfig.prepend(
@@ -79,6 +87,7 @@ module Switchman
         ::ActiveRecord::Relation.include(ActiveRecord::QueryMethods)
         ::ActiveRecord::Relation.prepend(GuardRail::Relation)
         ::ActiveRecord::Relation.prepend(ActiveRecord::Relation)
+        ::ActiveRecord::Relation.prepend(ActiveRecord::Relation::InsertUpsertAll) if ::Rails.version >= "7.2"
         ::ActiveRecord::Relation.include(ActiveRecord::SpawnMethods)
         ::ActiveRecord::Relation.include(CallSuper)
 

data/lib/switchman/shard.rb

@@ -5,8 +5,10 @@ module Switchman
     # ten trillion possible ids per shard. yup.
     IDS_PER_SHARD = 10_000_000_000_000
 
+    # rubocop:disable Style/SymbolProc -- transforming to a lambda produces "no receiver given"
     # only allow one default
     validates_uniqueness_of :default, if: ->(s) { s.default? }
+    # rubocop:enable Style/SymbolProc
 
     after_save :clear_cache
     after_destroy :clear_cache
@@ -51,7 +53,7 @@ module Switchman
       return [default] unless default.is_a?(Switchman::Shard)
       return all if !Switchman.region || DatabaseServer.none?(&:region)
 
-      in_region(Switchman.region, include_regionless: include_regionless)
+      in_region(Switchman.region, include_regionless:)
     end)
 
     class << self
@@ -141,7 +143,7 @@ module Switchman
 
         unless cached_shards.key?(id)
           cached_shards[id] = Shard.default.activate do
-            find_cached(["shard", id]) { find_by(id: id) }
+            find_cached(["shard", id]) { find_by(id:) }
           end
         end
         cached_shards[id]
@@ -233,7 +235,7 @@ module Switchman
               Switchman.config[:on_fork_proc]&.call
               with_each_shard(subscope,
                               classes,
-                              exception: exception,
+                              exception:,
                               output: output || :decorated) do
                 last_description = Shard.current.description
                 Parallel::ResultWrapper.new(yield)
@@ -459,7 +461,7 @@ module Switchman
             connects_to_hash.each do |(db_name, role_hash)|
               role_hash.each_key do |role|
                 role_hash.delete(role) if klass.connection_handler.retrieve_connection_pool(
-                  klass.connection_specification_name, role: role, shard: db_name
+                  klass.connection_specification_name, role:, shard: db_name
                 )
               end
             end
@@ -590,9 +592,9 @@ module Switchman
       id
     end
 
-    def activate(*classes, &block)
+    def activate(*classes, &)
       shards = hashify_classes(classes)
-      Shard.activate(shards, &block)
+      Shard.activate(shards, &)
     end
 
     # for use from console ONLY

data/lib/switchman/sharded_instrumenter.rb

@@ -19,7 +19,7 @@ module Switchman
           env: @shard_host.pool.connection_class&.current_role
         }
       end
-      super(name, payload)
+      super
     end
   end
 end

data/lib/switchman/test_helper.rb

@@ -63,7 +63,7 @@ module Switchman
 
       def find_existing_test_shard(server, name)
         if server == Shard.default.database_server
-          server.shards.where(name: name).first
+          server.shards.where(name:).first
         else
           shard = Shard.where("database_server_id IS NOT NULL AND name=?", name).first
           # if somehow databases got created in a different order, change the shard to match

data/lib/switchman/version.rb

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

data/lib/tasks/switchman.rake

@@ -94,7 +94,7 @@ module Switchman
                      else
                        nil
                      end
-            Shard.with_each_shard(scope, classes, output: output, **options) do
+            Shard.with_each_shard(scope, classes, output:, **options) do
               shard = Shard.current
 
               if log_format == "json"

metadata

@@ -1,16 +1,16 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 4.0.1
+  version: 4.1.0
 platform: ruby
 authors:
 - Cody Cutrer
 - James Williams
 - Jacob Fugal
-autorequire: 
+autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-02-18 00:00:00.000000000 Z
+date: 2025-03-21 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -21,7 +21,7 @@ dependencies:
         version: '7.0'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.2'
+        version: '7.3'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
@@ -31,7 +31,7 @@ dependencies:
         version: '7.0'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.2'
+        version: '7.3'
 - !ruby/object:Gem::Dependency
   name: guardrail
   requirement: !ruby/object:Gem::Requirement
@@ -69,7 +69,7 @@ dependencies:
         version: '7.0'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.2'
+        version: '7.3'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
@@ -79,7 +79,7 @@ dependencies:
         version: '7.0'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.2'
+        version: '7.3'
 - !ruby/object:Gem::Dependency
   name: debug
   requirement: !ruby/object:Gem::Requirement
@@ -198,14 +198,28 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.2'
+        version: '3.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.2'
+        version: '3.0'
+- !ruby/object:Gem::Dependency
+  name: rubocop-rspec_rails
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.29'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.29'
 - !ruby/object:Gem::Dependency
   name: simplecov
   requirement: !ruby/object:Gem::Requirement
@@ -290,7 +304,7 @@ licenses:
 metadata:
   rubygems_mfa_required: 'true'
   source_code_uri: https://github.com/instructure/switchman
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -298,15 +312,15 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '3.0'
+      version: '3.1'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.2.33
-signing_key: 
+rubygems_version: 3.3.27
+signing_key:
 specification_version: 4
 summary: Rails sharding magic
 test_files: []