switchman 3.3.6 → 4.2.5

data/lib/switchman/parallel.rb

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

data/lib/switchman/r_spec_helper.rb

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

data/lib/switchman/shard.rb

@@ -5,14 +5,57 @@ 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
 
     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:)
+    end)
+
     class << self
       def sharded_models
         @sharded_models ||= [::ActiveRecord::Base, UnshardedRecord].freeze
@@ -36,13 +79,15 @@ module Switchman
 
           # Now find the actual record, if it exists
           @default = begin
-            find_cached('default_shard') { Shard.where(default: true).take } || default
+            find_cached("default_shard") { Shard.where(default: true).take } || default
           rescue
             default
           end
 
           # make sure this is not erroneously cached
-          @default.database_server.remove_instance_variable(:@primary_shard) if @default.database_server.instance_variable_defined?(:@primary_shard)
+          if @default.database_server.instance_variable_defined?(:@primary_shard)
+            @default.database_server.remove_instance_variable(:@primary_shard)
+          end
 
           # and finally, check for cached references to the default shard on the existing connection
           sharded_models.each do |klass|
@@ -75,28 +120,30 @@ module Switchman
                       klass.current_switchman_shard != shard
 
           (activated_classes ||= []) << klass
-          klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass], switchman_shard: shard }
+          klass.connected_to_stack << { shard: shard.database_server.id.to_sym,
+                                        klasses: [klass],
+                                        switchman_shard: shard }
         end
         activated_classes
       end
 
       def active_shards
-        sharded_models.map do |klass|
+        sharded_models.filter_map do |klass|
           [klass, current(klass)]
-        end.compact.to_h
+        end.to_h
       end
 
       def lookup(id)
         id_i = id.to_i
-        return current if id_i == current.id || id == 'self'
-        return default if id_i == default.id || id.nil? || id == 'default'
+        return current if id_i == current.id || id == "self"
+        return default if id_i == default.id || id.nil? || id == "default"
 
         id = id_i
         raise ArgumentError if id.zero?
 
         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]
@@ -120,7 +167,7 @@ module Switchman
       #                forking.
       #    exception: - :ignore, :raise, :defer (wait until the end and raise the first
       #                error), or a proc
-      #    output: - :simple, :decorated (with database_server_id:shard_name)
+      #    output: - :simple, :decorated (with database_server_id:shard_name), custom lambda transformer
       def with_each_shard(*args, parallel: false, exception: :raise, output: :simple)
         raise ArgumentError, "wrong number of arguments (#{args.length} for 0...2)" if args.length > 2
 
@@ -136,19 +183,28 @@ module Switchman
           scope, classes = args
         end
 
+        output = if output == :decorated
+                   ->(arg) { "#{Shard.current.description}: #{arg}" }
+                 elsif output == :simple
+                   nil
+                 else
+                   output
+                 end
         parallel = [Environment.cpu_count || 2, 2].min if parallel == true
         parallel = 0 if parallel == false || parallel.nil?
 
         scope ||= Shard.all
-        scope = scope.order(::Arel.sql('database_server_id IS NOT NULL, database_server_id, id')) if ::ActiveRecord::Relation === scope && scope.order_values.empty?
+        if ::ActiveRecord::Relation === scope && scope.order_values.empty?
+          scope = scope.order(::Arel.sql("database_server_id IS NOT NULL, database_server_id, id"))
+        end
 
         if parallel > 1
           if ::ActiveRecord::Relation === scope
             # still need a post-uniq, cause the default database server could be NULL or Rails.env in the db
-            database_servers = scope.reorder('database_server_id').select(:database_server_id).distinct.
-                               map(&:database_server).compact.uniq
+            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)]
@@ -160,20 +216,23 @@ 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!
+          ::ActiveRecord::Base.connection_handler.clear_all_connections!(:all)
 
-          parent_process_name = `ps -ocommand= -p#{Process.pid}`.slice(/#{$0}.*/)
-          ret = ::Parallel.map(scopes, in_processes: scopes.length > 1 ? parallel : 0) do |server, subscope|
+          parent_process_name = sanitized_process_title
+          ret = ::Parallel.map(scopes, in_processes: (scopes.length > 1) ? parallel : 0) do |server, subscope|
             name = server.id
             last_description = name
 
             begin
               max_length = 128 - name.length - 3
               short_parent_name = parent_process_name[0..max_length] if max_length >= 0
-              new_title = [short_parent_name, name].join(' ')
+              new_title = [short_parent_name, name].join(" ")
               Process.setproctitle(new_title)
               Switchman.config[:on_fork_proc]&.call
-              with_each_shard(subscope, classes, exception: exception, output: :decorated) do
+              with_each_shard(subscope,
+                              classes,
+                              exception:,
+                              output: output || :decorated) do
                 last_description = Shard.current.description
                 Parallel::ResultWrapper.new(yield)
               end
@@ -187,8 +246,9 @@ module Switchman
           unless errors.empty?
             raise errors.first.exception if errors.length == 1
 
+            errors_desc = errors.map(&:name).sort.join(", ")
             raise Errors::ParallelShardExecError,
-                  "The following database server(s) did not finish processing cleanly: #{errors.map(&:name).sort.join(', ')}",
+                  "The following database server(s) did not finish processing cleanly: #{errors_desc}",
                   cause: errors.first.exception
           end
 
@@ -206,9 +266,9 @@ module Switchman
           next unless shard.database_server
 
           shard.activate(*classes) do
-            if output == :decorated
-              $stdout = Parallel::PrefixingIO.new(shard.description, $stdout)
-              $stderr = Parallel::PrefixingIO.new(shard.description, $stderr)
+            if output
+              $stdout = Parallel::TransformingIO.new(output, $stdout)
+              $stderr = Parallel::TransformingIO.new(output, $stderr)
             end
 
             result.concat Array.wrap(yield)
@@ -251,7 +311,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
@@ -293,14 +353,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
@@ -384,7 +444,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
@@ -397,13 +457,16 @@ 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
           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
 
@@ -437,8 +500,50 @@ module Switchman
         shard = nil unless shard.database_server
         shard
       end
+
+      # Determines the name of the current process, including arguments, but stripping
+      # any shebang from the invoked script, and any additional path info from the
+      # executable.
+      #
+      # @return [String]
+      def sanitized_process_title
+        # get the effective process name from `ps`; this will include any changes
+        # from Process.setproctitle _or_ assigning to $0.
+        parent_process_name = `ps -ocommand= -p#{Process.pid}`.strip
+        # Effective process titles may be shorter than the actual
+        # command; truncate our ARGV[0] so that they are comparable
+        # for the next step
+        argv0 = if parent_process_name.length < Process.argv0.length
+                  Process.argv0[0..parent_process_name.length]
+                else
+                  Process.argv0
+                end
+
+        # when running via a shebang, the `ps` output will include the shebang
+        # (i.e. it will be "ruby bin/rails c"); attempt to strip it off.
+        # Note that argv0 in this case will _only_ be `bin/rails` (no shebang,
+        # no arguments). We want to preserve the arguments we got from `ps`
+        if (index = parent_process_name.index(argv0))
+          parent_process_name.slice!(0...index)
+        end
+
+        # remove directories from the main executable to make more room
+        # for additional info
+        argv = parent_process_name.shellsplit
+        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
@@ -467,7 +572,7 @@ module Switchman
     end
 
     def description
-      [database_server.id, name].compact.join(':')
+      [database_server.id, name].compact.join(":")
     end
 
     # Shards are always on the default shard
@@ -479,9 +584,13 @@ module Switchman
       id
     end
 
-    def activate(*classes, &block)
+    def original_id_value
+      id
+    end
+
+    def activate(*classes, &)
       shards = hashify_classes(classes)
-      Shard.activate(shards, &block)
+      Shard.activate(shards, &)
     end
 
     # for use from console ONLY
@@ -501,48 +610,39 @@ module Switchman
     end
 
     def drop_database
-      raise('Cannot drop the database of the default shard') if default?
+      raise "Cannot drop the database of the default shard" if default?
       return unless read_attribute(:name)
 
       begin
-        adapter = database_server.config[:adapter]
-        sharding_config = Switchman.config || {}
-        drop_statement = sharding_config[adapter]&.[](:drop_statement)
-        drop_statement ||= sharding_config[:drop_statement]
-        if drop_statement
-          drop_statement = Array(drop_statement).dup.
-                           map { |statement| statement.gsub('%{name}', name) }
+        activate do
+          self.class.drop_database(name)
         end
+      rescue ::ActiveRecord::StatementInvalid => e
+        logger.error "Drop failed: #{e}"
+      end
+    end
 
-        case adapter
-        when 'mysql', 'mysql2'
-          activate do
-            ::GuardRail.activate(:deploy) do
-              drop_statement ||= "DROP DATABASE #{name}"
-              Array(drop_statement).each do |stmt|
-                ::ActiveRecord::Base.connection.execute(stmt)
-              end
-            end
-          end
-        when 'postgresql'
-          activate do
-            ::GuardRail.activate(:deploy) do
-              # Shut up, Postgres!
-              conn = ::ActiveRecord::Base.connection
-              old_proc = conn.raw_connection.set_notice_processor {}
-              begin
-                drop_statement ||= "DROP SCHEMA #{name} CASCADE"
-                Array(drop_statement).each do |stmt|
-                  ::ActiveRecord::Base.connection.execute(stmt)
-                end
-              ensure
-                conn.raw_connection.set_notice_processor(&old_proc) if old_proc
-              end
-            end
+    #
+    # Drops a specific database/schema from the currently active connection
+    #
+    def self.drop_database(name)
+      sharding_config = Switchman.config || {}
+      drop_statement = sharding_config["postgresql"]&.[](:drop_statement)
+      drop_statement ||= sharding_config[:drop_statement]
+      drop_statement = Array(drop_statement).map { |statement| statement.gsub("%{name}", name) } if drop_statement
+
+      ::GuardRail.activate(:deploy) do
+        # Shut up, Postgres!
+        conn = ::ActiveRecord::Base.connection
+        old_proc = conn.raw_connection.set_notice_processor {}
+        begin
+          drop_statement ||= "DROP SCHEMA #{name} CASCADE"
+          Array(drop_statement).each do |stmt|
+            ::ActiveRecord::Base.connection.execute(stmt)
           end
+        ensure
+          conn.raw_connection.set_notice_processor(&old_proc) if old_proc
         end
-      rescue
-        logger.info "Drop failed: #{$!}"
       end
     end
 
@@ -561,7 +661,7 @@ module Switchman
     end
 
     def destroy
-      raise('Cannot destroy the default shard') if default?
+      raise("Cannot destroy the default shard") if default?
 
       super
     end
@@ -570,8 +670,8 @@ module Switchman
 
     def clear_cache
       Shard.default.activate do
-        Switchman.cache.delete(['shard', id].join('/'))
-        Switchman.cache.delete('default_shard') if default?
+        Switchman.cache.delete(["shard", id].join("/"))
+        Switchman.cache.delete("default_shard") if default?
       end
       self.class.clear_cache
     end

data/lib/switchman/sharded_instrumenter.rb

@@ -3,7 +3,7 @@
 module Switchman
   class ShardedInstrumenter < ::SimpleDelegator
     def initialize(instrumenter, shard_host)
-      super instrumenter
+      super(instrumenter)
       @shard_host = shard_host
     end
 
@@ -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: ::Rails.version < '7.0' ? @shard_host.pool.connection_klass&.current_role : @shard_host.pool.connection_class&.current_role
+          env:
         }
       end
-      super name, payload
+      super
     end
   end
 end

data/lib/switchman/standard_error.rb

@@ -13,6 +13,10 @@ module Switchman
         Thread.current[:switchman_error_handler] = true
 
         @active_shards ||= Shard.active_shards
+      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
@@ -63,9 +63,9 @@ 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
+          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
           shard.database_server = server if shard
           shard

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = '3.3.6'
+  VERSION = "4.2.5"
 end