switchman 4.0.0 → 4.2.5

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
@@ -110,7 +163,7 @@ module Switchman
         end
       end
 
-      def activate(unordered: false, &block)
+      def activate(count: false, unordered: false, &block)
         shards = all_shards
         if Array === shards && shards.length == 1
           if !loaded? && shard_value != shards.first
@@ -136,7 +189,13 @@ module Switchman
 
             shard_results = relation.activate(&block)
 
-            if shard_results.present? && !unordered
+            if shard_results.present? && count
+              unless shard_results.is_a?(Integer)
+                raise "expected integer result for count, got #{shard_results.class.name}"
+              end
+
+              result_count += shard_results
+            elsif shard_results.present? && !unordered
               can_order ||= can_order_cross_shard_results? unless order_values.empty?
               raise OrderOnMultiShardQuery if !can_order && !order_values.empty? && result_count.positive?
 

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

@@ -29,14 +29,14 @@ module Switchman
         params, connection = args
         klass = @klass
         target_shard = nil
-        if (primary_index = bind_map.primary_value_index)
+        if (primary_index = @bind_map.primary_value_index)
           primary_value = params[primary_index]
           target_shard = Shard.local_id_for(primary_value)[1]
         end
         current_shard = Shard.current(klass.connection_class_for_self)
         target_shard ||= current_shard
 
-        bind_values = bind_map.bind(params, current_shard, target_shard)
+        bind_values = @bind_map.bind(params, current_shard, target_shard)
 
         target_shard.activate(klass.connection_class_for_self) do
           sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)

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,85 @@ 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 = (payload[:connection_name] if payload.key?(:connection_name))
+              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:)
+                  connection.connect! # 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:)
+              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/arel.rb

@@ -38,31 +38,6 @@ module Switchman
           collector << quote_local_table_name(join_name) << "." << quote_column_name(o.name)
         end
 
-        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 << if o.type == :in
-                           " IN ("
-                         else
-                           " NOT IN ("
-                         end
-
-            values = o.casted_values
-
-            if values.empty?
-              collector << @connection.quote(nil)
-            else
-              collector.add_binds(values, o.proc_for_binds, &bind_block)
-            end
-
-            collector << ")"
-            collector
-          end
-        end
-
         # rubocop:enable Naming/MethodName
         # rubocop:enable Naming/MethodParameterName
 

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(
@@ -259,11 +259,7 @@ module Switchman
           unless schema == false
             shard.activate do
               ::ActiveRecord::Base.connection.transaction(requires_new: true) do
-                if ::Rails.version < "7.1"
-                  ::ActiveRecord::Base.connection.migration_context.migrate
-                else
-                  ::ActiveRecord::MigrationContext.new(::ActiveRecord::Migrator.migrations_paths).migrate
-                end
+                ::ActiveRecord::MigrationContext.new(::ActiveRecord::Migrator.migrations_paths).migrate
               end
 
               ::ActiveRecord::Base.descendants.reject do |m|

data/lib/switchman/engine.rb

@@ -4,8 +4,6 @@ module Switchman
   class Engine < ::Rails::Engine
     isolate_namespace Switchman
 
-    # enable Rails 6.1 style connection handling
-    config.active_record.legacy_connection_handling = false if ::Rails.version < "7.1"
     config.active_record.writing_role = :primary
 
     ::GuardRail.singleton_class.prepend(GuardRail::ClassMethods)
@@ -47,12 +45,18 @@ module Switchman
           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::ConnectionHandler.prepend(ActiveRecord::ConnectionHandler)
         ::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 +83,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]
@@ -202,7 +204,7 @@ module Switchman
             database_servers = scope.reorder("database_server_id").select(:database_server_id).distinct
                                     .filter_map(&:database_server).uniq
             # nothing to do
-            return if database_servers.count.zero?
+            return if database_servers.none?
 
             scopes = database_servers.to_h do |server|
               [server, scope.merge(server.shards)]
@@ -214,11 +216,7 @@ 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)
-          if ::Rails.version < "7.1"
-            ::ActiveRecord::Base.clear_all_connections!(nil)
-          else
-            ::ActiveRecord::Base.connection_handler.clear_all_connections!(:all)
-          end
+          ::ActiveRecord::Base.connection_handler.clear_all_connections!(:all)
 
           parent_process_name = sanitized_process_title
           ret = ::Parallel.map(scopes, in_processes: (scopes.length > 1) ? parallel : 0) do |server, subscope|
@@ -233,7 +231,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 +457,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 +588,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

@@ -13,13 +13,19 @@ module Switchman
       # when we might be doing a query while defining attribute methods,
       # so just avoid logging then
       if shard.is_a?(Shard) && Shard.instance_variable_get(:@attribute_methods_generated)
+        env = if ::Rails.version < "8.0"
+                @shard_host.pool.connection_class&.current_role
+              else
+                @shard_host.pool.connection_descriptor&.name&.constantize&.current_role
+              end
+
         payload[:shard] = {
           database_server_id: shard.database_server.id,
           id: shard.id,
-          env: @shard_host.pool.connection_class&.current_role
+          env:
         }
       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.0"
+  VERSION = "4.2.5"
 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"