switchman 3.0.2 → 3.0.9

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1224320ec4f3ff0d8f05d1c5537ce06f33932e4c9da069d82183aa7562e1cab4
-  data.tar.gz: 1cd6f52c82fbdea9c3aac4c569b291428d2476abc4a2198c09beb1da64700473
+  metadata.gz: a8f502f6ec406e1292800c9ae65cbb75e584f1059c466ec732c8aeb50379c1df
+  data.tar.gz: 823fd74a80bf27e0ddfd29d22a03f21f0b76e3d119005211414d4e998f59d5ad
 SHA512:
-  metadata.gz: 052e1675cb5d6b07844dab24c802c6f98b8e75fe4c0146a407e275d89633b5c3058801538a2537f82a7798260ad9d5dd2989536ad88d06a37bb33cf9082444dd
-  data.tar.gz: ebaa91ea8f6be5f878ac2f8b57a7954853603cb20770f083263fd369774cedc4b33c8127a8b17746bbf4e5b91e1e742df2970be8dab08aa708baae3bc130b988
+  metadata.gz: f6010a838cbb0f17f74aa12a6859d41c3a6a44c0752b130077d9c762c1c4e7d37e32f0d06fb11aa26d6727b41e41d2b1873d1ac5f5889dc77815379148e332df
+  data.tar.gz: 7479c21a001fa920c4a81a9fb797a8468dd69dca4d190e636b9ceaf159381cf122c7e10395f3ecf4902bbd435841a318636375ae72130d7dfd8f9d4632d02002

data/app/models/switchman/shard.rb

@@ -39,12 +39,15 @@ module Switchman
           # the first time we need a dummy dummy for re-entrancy to avoid looping on ourselves
           @default ||= default
 
-          # Now find the actual record, if it exists; rescue the fake default if the table doesn't exist
+          # Now find the actual record, if it exists
           @default = begin
             find_cached('default_shard') { Shard.where(default: true).take } || default
+          # If we are *super* early in boot, the connection pool won't exist; we don't want to fill in the default shard yet
+          # Otherwise, rescue the fake default if the table doesn't exist
           rescue
-            default
+            sharding_initialized ? default : nil
           end
+          return default unless @default
 
           # make sure this is not erroneously cached
           @default.database_server.remove_instance_variable(:@primary_shard) if @default.database_server.instance_variable_defined?(:@primary_shard)
@@ -59,34 +62,38 @@ module Switchman
 
       def current(klass = ::ActiveRecord::Base)
         klass ||= ::ActiveRecord::Base
-        klass.connection_pool.shard
+        klass.current_switchman_shard
       end
 
       def activate(shards)
         activated_classes = activate!(shards)
         yield
       ensure
-        activated_classes.each do |klass|
-          klass.connection_pool.shard_stack.pop
+        activated_classes&.each do |klass|
           klass.connected_to_stack.pop
         end
       end
 
       def activate!(shards)
-        activated_classes = []
+        activated_classes = nil
         shards.each do |klass, shard|
           next if klass == UnshardedRecord
 
           next unless klass.current_shard != shard.database_server.id.to_sym ||
-                      klass.connection_pool.shard != shard
+                      klass.current_switchman_shard != shard
 
-          activated_classes << klass
-          klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass] }
-          klass.connection_pool.shard_stack << shard
+          (activated_classes ||= []) << klass
+          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|
+          [klass, current(klass)]
+        end.compact.to_h
+      end
+
       def lookup(id)
         id_i = id.to_i
         return current if id_i == current.id || id == 'self'
@@ -103,6 +110,11 @@ module Switchman
         cached_shards[id]
       end
 
+      def preload_cache
+        cached_shards.reverse_merge!(active_shards.values.index_by(&:id))
+        cached_shards.reverse_merge!(all.index_by(&:id))
+      end
+
       def clear_cache
         cached_shards.clear
       end
@@ -111,14 +123,13 @@ module Switchman
       #
       # * +shards+ - an array or relation of Shards to iterate over
       # * +classes+ - an array of classes to activate
-      #    parallel: - true/false to execute in parallel, or a integer of how many
-      #                sub-processes per database server. Note that parallel
-      #                invocation currently uses forking, so should be used sparingly
-      #                because errors are not raised, and you cannot get results back
-      #    max_procs: - only run this many parallel processes at a time
+      #    parallel: - true/false to execute in parallel, or an integer of how many
+      #                sub-processes. Note that parallel invocation currently uses
+      #                forking, so should be used sparingly because you cannot get
+      #                results back
       #    exception: - :ignore, :raise, :defer (wait until the end and raise the first
       #                error), or a proc
-      def with_each_shard(*args, parallel: false, max_procs: nil, exception: :raise, &block)
+      def with_each_shard(*args, parallel: false, exception: :raise, &block)
         raise ArgumentError, "wrong number of arguments (#{args.length} for 0...2)" if args.length > 2
 
         return Array.wrap(yield) unless default.is_a?(Shard)
@@ -133,14 +144,13 @@ module Switchman
           scope, classes = args
         end
 
-        parallel = 1 if parallel == true
+        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 parallel.positive?
-          max_procs = determine_max_procs(max_procs, parallel)
+        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.
@@ -148,39 +158,11 @@ module Switchman
             # nothing to do
             return if database_servers.count.zero?
 
-            parallel = [(max_procs.to_f / database_servers.count).ceil, parallel].min if max_procs
-
-            scopes = database_servers.map do |server|
-              server_scope = server.shards.merge(scope)
-              if parallel == 1
-                subscopes = [server_scope]
-              else
-                subscopes = []
-                total = server_scope.count
-                ranges = []
-                server_scope.find_ids_in_ranges(batch_size: (total.to_f / parallel).ceil) do |min, max|
-                  ranges << [min, max]
-                end
-                # create a half-open range on the last one
-                ranges.last[1] = nil
-                ranges.each do |min, max|
-                  subscope = server_scope.where('id>=?', min)
-                  subscope = subscope.where('id<=?', max) if max
-                  subscopes << subscope
-                end
-              end
-              [server, subscopes]
-            end.to_h
+            scopes = database_servers.to_h do |server|
+              [server, server.shards.merge(scope)]
+            end
           else
             scopes = scope.group_by(&:database_server)
-            if parallel > 1
-              parallel = [(max_procs.to_f / scopes.count).ceil, parallel].min if max_procs
-              scopes = scopes.map do |(server, shards)|
-                [server, shards.in_groups(parallel, false).compact]
-              end.to_h
-            else
-              scopes = scopes.map { |(server, shards)| [server, [shards]] }.to_h
-            end
           end
 
           exception_pipes = []
@@ -206,80 +188,83 @@ module Switchman
           end
 
           # only one process; don't bother forking
-          if scopes.length == 1 && parallel == 1
-            return with_each_shard(scopes.first.last.first, classes, exception: exception,
-                                   &block)
-          end
+          return with_each_shard(scopes.first.last, classes, exception: exception, &block) if scopes.length == 1
 
           # 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!
 
-          scopes.each do |server, subscopes|
-            subscopes.each_with_index do |subscope, idx|
-              name = if subscopes.length > 1
-                       "#{server.id} #{idx + 1}"
-                     else
-                       server.id
-                     end
-
-              exception_pipe = IO.pipe
-              exception_pipes << exception_pipe
-              pid, io_in, io_out, io_err = Open4.pfork4(lambda do
-                Switchman.config[:on_fork_proc]&.call
-
-                # set a pretty name for the process title, up to 128 characters
-                # (we don't actually know the limit, depending on how the process
-                # was started)
-                # first, simplify the binary name by stripping directories,
-                # then truncate arguments as necessary
-                bin = File.basename($0) # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
-                max_length = 128 - bin.length - name.length - 3
-                args = ARGV.join(' ')
-                args = args[0..max_length] if max_length >= 0
-                new_title = [bin, args, name].join(' ')
-                Process.setproctitle(new_title)
-
-                with_each_shard(subscope, classes, exception: exception, &block)
-                exception_pipe.last.close
-              rescue => e
-                begin
-                  dumped = Marshal.dump(e)
-                rescue
-                  # couldn't dump the exception; create a copy with just
-                  # the message and the backtrace
-                  e2 = e.class.new(e.message)
-                  e2.set_backtrace(e.backtrace)
-                  e2.instance_variable_set(:@active_shards, e.instance_variable_get(:@active_shards))
-                  dumped = Marshal.dump(e2)
-                end
-                exception_pipe.last.set_encoding(dumped.encoding)
-                exception_pipe.last.write(dumped)
-                exception_pipe.last.flush
-                exception_pipe.last.close
-                exit! 1
-              end)
+          scopes.each do |server, subscope|
+            name = server.id
+
+            exception_pipe = IO.pipe
+            exception_pipes << exception_pipe
+            pid, io_in, io_out, io_err = Open4.pfork4(lambda do
+              Switchman.config[:on_fork_proc]&.call
+
+              # set a pretty name for the process title, up to 128 characters
+              # (we don't actually know the limit, depending on how the process
+              # was started)
+              # first, simplify the binary name by stripping directories,
+              # then truncate arguments as necessary
+              bin = File.basename($0) # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
+              max_length = 128 - bin.length - name.length - 3
+              args = ARGV.join(' ')
+              args = args[0..max_length] if max_length >= 0
+              new_title = [bin, args, name].join(' ')
+              Process.setproctitle(new_title)
+
+              with_each_shard(subscope, classes, exception: exception, &block)
               exception_pipe.last.close
-              pids << pid
-              io_in.close # don't care about writing to stdin
-              out_fds << io_out
-              err_fds << io_err
-              pid_to_name_map[pid] = name
-              fd_to_name_map[io_out] = name
-              fd_to_name_map[io_err] = name
-
-              while max_procs && pids.count >= max_procs
-                while max_procs && out_fds.count >= max_procs
-                  # wait for output if we've hit the max_procs limit
-                  wait_for_output.call
-                end
-                # we've gotten all the output from one fd so wait for its child process to exit
-                found_pid, status = Process.wait2
-                pids.delete(found_pid)
-                errors << pid_to_name_map[found_pid] if status.exitstatus != 0
+            rescue Exception => e # rubocop:disable Lint/RescueException
+              begin
+                dumped = Marshal.dump(e)
+                dumped = nil if dumped.length > 64 * 1024
+              rescue
+                dumped = nil
               end
+
+              if dumped.nil?
+                # couldn't dump the exception; create a copy with just
+                # the message and the backtrace
+                e2 = e.class.new(e.message)
+                backtrace = e.backtrace
+                # truncate excessively long backtraces
+                backtrace = backtrace[0...25] + ['...'] + backtrace[-25..] if backtrace.length > 50
+                e2.set_backtrace(backtrace)
+                e2.instance_variable_set(:@active_shards, e.instance_variable_get(:@active_shards))
+                dumped = Marshal.dump(e2)
+              end
+              exception_pipe.last.set_encoding(dumped.encoding)
+              exception_pipe.last.write(dumped)
+              exception_pipe.last.flush
+              exception_pipe.last.close
+              exit! 1
+            end)
+            exception_pipe.last.close
+            pids << pid
+            io_in.close # don't care about writing to stdin
+            out_fds << io_out
+            err_fds << io_err
+            pid_to_name_map[pid] = name
+            fd_to_name_map[io_out] = name
+            fd_to_name_map[io_err] = name
+
+            while pids.count >= parallel
+              while out_fds.count >= parallel
+                # wait for output if we've hit the parallel limit
+                wait_for_output.call
+              end
+              # we've gotten all the output from one fd so wait for its child process to exit
+              found_pid, status = Process.wait2
+              pids.delete(found_pid)
+              errors << pid_to_name_map[found_pid] if status.exitstatus != 0
             end
+            # we've gotten all the output from one fd so wait for its child process to exit
+            found_pid, status = Process.wait2
+            pids.delete(found_pid)
+            errors << pid_to_name_map[found_pid] if status.exitstatus != 0
           end
 
           wait_for_output.call while out_fds.any? || err_fds.any?
@@ -402,7 +387,7 @@ module Switchman
           signed_id_operation(local_id) do |id|
             return nil if id > IDS_PER_SHARD
 
-            $1.to_i * IDS_PER_SHARD + id
+            ($1.to_i * IDS_PER_SHARD) + id
           end
         when Integer, /^-?\d+$/
           any_id.to_i
@@ -480,32 +465,19 @@ module Switchman
         shard || source_shard || Shard.current
       end
 
-      # given the provided option, determines whether we need to (and whether
-      # it's possible) to determine a reasonable default.
-      def determine_max_procs(max_procs_input, parallel_input = 2)
-        max_procs = nil
-        if max_procs_input
-          max_procs = max_procs_input.to_i
-          max_procs = nil if max_procs.zero?
-        else
-          return 1 if parallel_input.nil? || parallel_input < 1
-
-          cpus = Environment.cpu_count
-          max_procs = cpus * parallel_input if cpus&.positive?
-        end
+      private
 
-        max_procs
+      def sharding_initialized
+        @sharding_initialized ||= false
       end
 
-      private
-
       def add_sharded_model(klass)
         @sharded_models = (sharded_models + [klass]).freeze
         initialize_sharding
       end
 
       def initialize_sharding
-        full_connects_to_hash = DatabaseServer.all.map { |db| [db.id.to_sym, db.connects_to_hash] }.to_h
+        full_connects_to_hash = DatabaseServer.all.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
@@ -526,6 +498,8 @@ module Switchman
 
           klass.connects_to shards: connects_to_hash
         end
+
+        @sharding_initialized = true
       end
 
       # in-process caching
@@ -596,6 +570,10 @@ module Switchman
       Shard.default
     end
 
+    def original_id
+      id
+    end
+
     def activate(*classes, &block)
       shards = hashify_classes(classes)
       Shard.activate(shards, &block)
@@ -668,7 +646,7 @@ module Switchman
       return nil unless local_id
 
       self.class.signed_id_operation(local_id) do |abs_id|
-        abs_id + id * IDS_PER_SHARD
+        abs_id + (id * IDS_PER_SHARD)
       end
     end
 

data/db/migrate/20180828183945_add_default_shard_index.rb

@@ -2,7 +2,7 @@
 
 class AddDefaultShardIndex < ActiveRecord::Migration[4.2]
   def change
-    Switchman::Shard.where(default: nil).update_all(default: false)
+    Switchman::Shard.where(default: nil).update_all(default: false) if Switchman::Shard.current.default?
     change_column_default :switchman_shards, :default, false
     change_column_null :switchman_shards, :default, false
     options = if connection.adapter_name == 'PostgreSQL'

data/db/migrate/20180828192111_add_timestamps_to_shards.rb

@@ -4,7 +4,7 @@ class AddTimestampsToShards < ActiveRecord::Migration[4.2]
   disable_ddl_transaction!
 
   def change
-    add_timestamps :switchman_shards, null: true
+    add_timestamps :switchman_shards, null: true, if_not_exists: true
     now = Time.now.utc
     Switchman::Shard.update_all(updated_at: now, created_at: now) if Switchman::Shard.current.default?
     change_column_null :switchman_shards, :updated_at, false

data/lib/switchman/active_record/abstract_adapter.rb

@@ -40,15 +40,6 @@ module Switchman
       ensure
         @last_query_at = Time.now
       end
-
-      private
-
-      def id_value_for_database(value)
-        return super unless value.class.sharded_primary_key?
-
-        # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
-        quote(value.id)
-      end
     end
   end
 end

data/lib/switchman/active_record/attribute_methods.rb

@@ -40,7 +40,11 @@ module Switchman
           if sharded_column?(attr_name)
             owner << <<-RUBY
               def global_#{attr_name}
-                ::Switchman::Shard.global_id_for(original_#{attr_name}, shard)
+                raw_value = original_#{attr_name}
+                return nil if raw_value.nil?
+                return raw_value if raw_value > ::Switchman::Shard::IDS_PER_SHARD
+
+                ::Switchman::Shard.global_id_for(raw_value, shard)
               end
             RUBY
           else
@@ -52,7 +56,9 @@ module Switchman
           if sharded_column?(attr_name)
             owner << <<-RUBY
               def local_#{attr_name}
-                ::Switchman::Shard.local_id_for(original_#{attr_name}).first
+                raw_value = original_#{attr_name}
+                return nil if raw_value.nil?
+                return raw_value % ::Switchman::Shard::IDS_PER_SHARD
               end
             RUBY
           else
@@ -104,7 +110,22 @@ module Switchman
               alias_method 'original_#{attr_name}', '#{attr_name}'
               # and replace with one that transposes the id
               def #{attr_name}
-                ::Switchman::Shard.relative_id_for(original_#{attr_name}, shard, ::Switchman::Shard.current(#{connection_classes_code_for_reflection(reflection)}))
+                raw_value = original_#{attr_name}
+                return nil if raw_value.nil?
+
+                abs_raw_value = raw_value.abs
+                current_shard = ::Switchman::Shard.current(#{connection_classes_code_for_reflection(reflection)})
+                same_shard = shard == current_shard
+                return raw_value if same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
+
+                value_shard_id = abs_raw_value / ::Switchman::Shard::IDS_PER_SHARD
+                # this is a stupid case when someone stuffed a global id for the current shard in instead
+                # of a local id
+                return raw_value % ::Switchman::Shard::IDS_PER_SHARD if value_shard_id == current_shard.id
+                return raw_value if !same_shard && abs_raw_value > ::Switchman::Shard::IDS_PER_SHARD
+                return shard.global_id_for(raw_value) if !same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
+
+                ::Switchman::Shard.relative_id_for(raw_value, shard, current_shard)
               end
 
               alias_method 'original_#{attr_name}=', '#{attr_name}='

data/lib/switchman/active_record/base.rb

@@ -77,17 +77,27 @@ module Switchman
 
           default_shard
         end
+
+        def current_switchman_shard
+          connected_to_stack.reverse_each do |hash|
+            return hash[:switchman_shard] if hash[:switchman_shard] && hash[:klasses].include?(connection_classes)
+          end
+
+          Shard.default
+        end
       end
 
-      def self.included(klass)
+      def self.prepended(klass)
         klass.singleton_class.prepend(ClassMethods)
-        klass.set_callback(:initialize, :before) do
-          @shard ||= if self.class.sharded_primary_key?
-                       Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
-                     else
-                       Shard.current(self.class.connection_classes)
-                     end
-        end
+      end
+
+      def _run_initialize_callbacks
+        @shard ||= if self.class.sharded_primary_key?
+                     Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
+                   else
+                     Shard.current(self.class.connection_classes)
+                   end
+        super
       end
 
       def shard
@@ -107,12 +117,12 @@ module Switchman
 
       def save(*, **)
         @shard_set_in_stone = true
-        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
+        super
       end
 
       def save!(*, **)
         @shard_set_in_stone = true
-        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
+        super
       end
 
       def destroy
@@ -134,6 +144,12 @@ module Switchman
         end
       end
 
+      def with_transaction_returning_status
+        shard.activate(self.class.connection_classes) do
+          super
+        end
+      end
+
       def hash
         self.class.sharded_primary_key? ? self.class.hash ^ global_id.hash : super
       end
@@ -149,19 +165,20 @@ module Switchman
         copy
       end
 
-      def quoted_id
-        return super unless self.class.sharded_primary_key?
+      def update_columns(*)
+        db = shard.database_server
+        return db.unguard { super } if ::GuardRail.environment != db.guard_rail_environment
 
-        # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
-        self.class.connection.quote(id)
+        super
       end
 
-      def update_columns(*)
-        db = Shard.current(self.class.connection_classes).database_server
-        if ::GuardRail.environment == db.guard_rail_environment
-          super
+      def id_for_database
+        if self.class.sharded_primary_key?
+          # It's an int, so so it's safe to just return it without passing it through anything else
+          # In theory we should do `@attributes[@primary_key].type.serialize(id)`, but that seems to have surprising side-effects
+          id
         else
-          db.unguard { super }
+          super
         end
       end
 
@@ -172,7 +189,7 @@ module Switchman
         if reflection
           if reflection.options[:polymorphic]
             begin
-              read_attribute(reflection.foreign_type)&.constantize&.connection_classes
+              read_attribute(reflection.foreign_type)&.constantize&.connection_classes || ::ActiveRecord::Base
             rescue NameError
               # in case someone is abusing foreign_type to not point to an actual class
               ::ActiveRecord::Base
@@ -182,7 +199,7 @@ module Switchman
             reflection.klass.connection_classes
           end
         else
-          connection_classes
+          self.class.connection_classes
         end
       end
     end

data/lib/switchman/active_record/calculations.rb

@@ -83,7 +83,7 @@ module Switchman
           if opts[:association]
             key_ids     = calculated_data.collect { |row| row[opts[:group_aliases].first] }
             key_records = opts[:association].klass.base_class.where(id: key_ids)
-            key_records = key_records.map { |r| [Shard.relative_id_for(r, shard, target_shard), r] }.to_h
+            key_records = key_records.to_h { |r| [Shard.relative_id_for(r, shard, target_shard), r] }
           end
 
           calculated_data.map do |row|

data/lib/switchman/active_record/connection_pool.rb

@@ -5,18 +5,6 @@ require 'switchman/errors'
 module Switchman
   module ActiveRecord
     module ConnectionPool
-      def shard
-        shard_stack.last || Shard.default
-      end
-
-      def shard_stack
-        unless (shard_stack = Thread.current.thread_variable_get(tls_key))
-          shard_stack = Concurrent::Array.new
-          Thread.current.thread_variable_set(tls_key, shard_stack)
-        end
-        shard_stack
-      end
-
       def default_schema
         connection unless @schemas
         # default shard will not switch databases immediately, so it won't be set yet
@@ -26,15 +14,15 @@ module Switchman
 
       def checkout_new_connection
         conn = super
-        conn.shard = shard
+        conn.shard = current_shard
         conn
       end
 
       def connection(switch_shard: true)
         conn = super()
-        raise NonExistentShardError if shard.new_record?
+        raise NonExistentShardError if current_shard.new_record?
 
-        switch_database(conn) if conn.shard != shard && switch_shard
+        switch_database(conn) if conn.shard != current_shard && switch_shard
         conn
       end
 
@@ -44,26 +32,18 @@ module Switchman
         flush
       end
 
-      def remove_shard!(shard)
-        synchronize do
-          # The shard might be currently active, so we need to update our own shard
-          self.shard = Shard.default if self.shard == shard
-          # Update out any connections that may be using this shard
-          @connections.each do |conn|
-            # This will also update the connection's shard to the default shard
-            switch_database(conn) if conn.shard == shard
-          end
-        end
-      end
-
       def switch_database(conn)
-        @schemas = conn.current_schemas if !@schemas && conn.adapter_name == 'PostgreSQL' && !shard.database_server.config[:shard_name]
+        @schemas = conn.current_schemas if !@schemas && conn.adapter_name == 'PostgreSQL' && !current_shard.database_server.config[:shard_name]
 
-        conn.shard = shard
+        conn.shard = current_shard
       end
 
       private
 
+      def current_shard
+        connection_klass.current_switchman_shard
+      end
+
       def tls_key
         "#{object_id}_shard".to_sym
       end

data/lib/switchman/active_record/migration.rb

@@ -14,16 +14,19 @@ module Switchman
 
       def connection
         conn = super
-        ::ActiveRecord::Base.connection_pool.switch_database(conn) if conn.shard != ::ActiveRecord::Base.connection_pool.shard
+        ::ActiveRecord::Base.connection_pool.switch_database(conn) if conn.shard != ::ActiveRecord::Base.current_switchman_shard
         conn
       end
     end
 
     module Migrator
-      # significant change: hash shard name, not database name
+      # significant change: just return MIGRATOR_SALT directly
+      # especially if you're going through pgbouncer, the database
+      # name you're accessing may not be consistent. it is NOT allowed
+      # to run migrations against multiple shards in the same database
+      # concurrently
       def generate_migrator_advisory_lock_id
-        shard_name_hash = Zlib.crc32(Shard.current.name)
-        ::ActiveRecord::Migrator::MIGRATOR_SALT * shard_name_hash
+        ::ActiveRecord::Migrator::MIGRATOR_SALT
       end
 
       # significant change: strip out prefer_secondary from config

data/lib/switchman/active_record/persistence.rb

@@ -11,6 +11,13 @@ module Switchman
       def update_columns(*)
         shard.activate(self.class.connection_classes) { super }
       end
+
+      def delete
+        db = shard.database_server
+        return db.unguard { super } unless ::GuardRail.environment == db.guard_rail_environment
+
+        super
+      end
     end
   end
 end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -32,7 +32,7 @@ module Switchman
       end
 
       # copy/paste; use quote_local_table_name
-      def drop_database(name) #:nodoc:
+      def drop_database(name) # :nodoc:
         execute "DROP DATABASE IF EXISTS #{quote_local_table_name(name)}"
       end
 
@@ -87,6 +87,10 @@ module Switchman
         name.quoted
       end
 
+      def with_global_table_name(&block)
+        with_local_table_name(false, &block)
+      end
+
       def with_local_table_name(enable = true) # rubocop:disable Style/OptionalBooleanParameter
         old_value = @use_local_table_name
         @use_local_table_name = enable

data/lib/switchman/active_record/predicate_builder.rb

@@ -15,7 +15,7 @@ module Switchman
         def convert_to_id(value)
           case value
           when ::ActiveRecord::Base
-            value.send(primary_key)
+            value.id
           else
             super
           end

data/lib/switchman/active_record/query_methods.rb

@@ -239,6 +239,10 @@ module Switchman
         connection.with_local_table_name { super }
       end
 
+      def table_name_matches?(from)
+        connection.with_global_table_name { super }
+      end
+
       def transpose_predicates(predicates,
                                source_shard,
                                target_shard,
@@ -266,11 +270,11 @@ module Switchman
           right_node = or_expr.right
           new_left_predicates = transpose_single_predicate(left_node, source_shard,
                                                            target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
-          or_expr.instance_variable_set(:@left, new_left_predicates) if new_left_predicates != left_node
           new_right_predicates = transpose_single_predicate(right_node, source_shard,
                                                             target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
-          or_expr.instance_variable_set(:@right, new_right_predicates) if new_right_predicates != right_node
-          return predicate
+          return predicate if new_left_predicates == left_node && new_right_predicates == right_node
+
+          return ::Arel::Nodes::Grouping.new ::Arel::Nodes::Or.new(new_left_predicates, new_right_predicates)
         end
         return predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
         return predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)

data/lib/switchman/active_record/relation.rb

@@ -63,7 +63,7 @@ module Switchman
       %I[update_all delete_all].each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
           def #{method}(*args)
-            result = self.activate { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
+            result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
             result = result.sum if result.is_a?(Array)
             result
           end
@@ -94,7 +94,7 @@ module Switchman
 
         while ids.first.present?
           ids.map!(&:to_i) if is_integer
-          ids << ids.first + batch_size if loose_mode
+          ids << (ids.first + batch_size) if loose_mode
 
           yield(*ids)
           last_value = ids.last
@@ -103,7 +103,7 @@ module Switchman
         end
       end
 
-      def activate(&block)
+      def activate(unordered: false, &block)
         shards = all_shards
         if Array === shards && shards.length == 1
           if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_classes)
@@ -124,7 +124,7 @@ module Switchman
 
             shard_results = relation.activate(&block)
 
-            if shard_results.present?
+            if 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?
 
@@ -148,8 +148,12 @@ module Switchman
         results.sort! do |l, r|
           result = 0
           order_values.each do |ov|
-            a = l.attribute(ov.expr.name)
-            b = r.attribute(ov.expr.name)
+            if l.is_a?(::ActiveRecord::Base)
+              a = l.attribute(ov.expr.name)
+              b = r.attribute(ov.expr.name)
+            else
+              a, b = l, r
+            end
             next if a == b
 
             if a.nil? || b.nil?

data/lib/switchman/active_record/test_fixtures.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+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 = payload[:spec_name] if payload.key?(:spec_name)
+          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)
+            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
+            end
+          end
+        end
+      end
+
+      def enlist_fixture_connections
+        setup_shared_connection_pool
+
+        ::ActiveRecord::Base.connection_handler.connection_pool_list.reject { |cp| FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym) }.map(&:connection)
+      end
+    end
+  end
+end

data/lib/switchman/arel.rb

@@ -11,22 +11,44 @@ module Switchman
     module Visitors
       module ToSql
         # rubocop:disable Naming/MethodName
+        # rubocop:disable Naming/MethodParameterName
 
-        def visit_Arel_Nodes_TableAlias(*args)
-          o, collector = args
+        def visit_Arel_Nodes_TableAlias(o, collector)
           collector = visit o.relation, collector
           collector << ' '
           collector << quote_local_table_name(o.name)
         end
 
-        def visit_Arel_Attributes_Attribute(*args)
-          o = args.first
+        def visit_Arel_Attributes_Attribute(o, collector)
           join_name = o.relation.table_alias || o.relation.name
-          result = "#{quote_local_table_name join_name}.#{quote_column_name o.name}"
-          args.last << result
+          collector << quote_local_table_name(join_name) << '.' << quote_column_name(o.name)
+        end
+
+        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
 
         # rubocop:enable Naming/MethodName
+        # rubocop:enable Naming/MethodParameterName
 
         def quote_local_table_name(name)
           return name if ::Arel::Nodes::SqlLiteral === name

data/lib/switchman/database_server.rb

@@ -59,7 +59,7 @@ module Switchman
       end
 
       def database_servers
-        unless @database_servers
+        if !@database_servers || @database_servers.empty?
           @database_servers = {}.with_indifferent_access
           ::ActiveRecord::Base.configurations.configurations.each do |config|
             if config.name.include?('/')
@@ -89,13 +89,13 @@ module Switchman
     end
 
     def connects_to_hash
-      self.class.all_roles.map do |role|
+      self.class.all_roles.to_h do |role|
         config_role = role
         config_role = :primary unless roles.include?(role)
         config_name = :"#{id}/#{config_role}"
         config_name = :primary if id == ::Rails.env && config_role == :primary
         [role.to_sym, config_name]
-      end.to_h
+      end
     end
 
     def destroy
@@ -186,58 +186,56 @@ module Switchman
 
       name ||= "#{config[:database]}_shard_#{id}"
 
+      schema_already_existed = false
+      shard = nil
       Shard.connection.transaction do
-        shard = Shard.create!(id: id,
-                              name: name,
-                              database_server_id: self.id)
-        schema_already_existed = false
-
-        begin
-          self.class.creating_new_shard = true
-          DatabaseServer.send(:reference_role, :deploy)
-          ::ActiveRecord::Base.connected_to(shard: self.id.to_sym, role: :deploy) do
-            if create_statement
-              if ::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}")
-                schema_already_existed = true
-                raise 'This schema already exists; cannot overwrite'
-              end
-              Array(create_statement.call).each do |stmt|
-                ::ActiveRecord::Base.connection.execute(stmt)
-              end
+        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,
+                                database_server_id: self.id)
+          if create_statement
+            if ::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}")
+              schema_already_existed = true
+              raise 'This schema already exists; cannot overwrite'
             end
-            if config[:adapter] == 'postgresql'
-              old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor do
-              end
+            Array(create_statement.call).each do |stmt|
+              ::ActiveRecord::Base.connection.execute(stmt)
             end
-            old_verbose = ::ActiveRecord::Migration.verbose
-            ::ActiveRecord::Migration.verbose = false
-
-            unless schema == false
-              shard.activate(*Shard.sharded_models) do
-                reset_column_information
-
-                ::ActiveRecord::Base.connection.transaction(requires_new: true) do
-                  ::ActiveRecord::Base.connection.migration_context.migrate
-                end
-                reset_column_information
-                ::ActiveRecord::Base.descendants.reject do |m|
-                  m <= UnshardedRecord || !m.table_exists?
-                end.each(&:define_attribute_methods)
+          end
+          if config[:adapter] == 'postgresql'
+            old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor do
+            end
+          end
+          old_verbose = ::ActiveRecord::Migration.verbose
+          ::ActiveRecord::Migration.verbose = false
+
+          unless schema == false
+            shard.activate do
+              reset_column_information
+
+              ::ActiveRecord::Base.connection.transaction(requires_new: true) do
+                ::ActiveRecord::Base.connection.migration_context.migrate
               end
+              reset_column_information
+              ::ActiveRecord::Base.descendants.reject do |m|
+                m <= UnshardedRecord || !m.table_exists?
+              end.each(&:define_attribute_methods)
             end
-          ensure
-            ::ActiveRecord::Migration.verbose = old_verbose
-            ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
           end
-          shard
-        rescue
-          shard.destroy
-          shard.drop_database rescue nil unless schema_already_existed
-          reset_column_information unless schema == false rescue nil
-          raise
         ensure
-          self.class.creating_new_shard = false
+          ::ActiveRecord::Migration.verbose = old_verbose
+          ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
         end
+        shard
+      rescue
+        shard&.destroy
+        shard&.drop_database rescue nil unless schema_already_existed
+        reset_column_information unless schema == false rescue nil
+        raise
+      ensure
+        self.class.creating_new_shard = false
       end
     end
 

data/lib/switchman/engine.rb

@@ -90,6 +90,7 @@ module Switchman
         require 'switchman/active_record/statement_cache'
         require 'switchman/active_record/tasks/database_tasks'
         require 'switchman/active_record/type_caster'
+        require 'switchman/active_record/test_fixtures'
         require 'switchman/arel'
         require 'switchman/call_super'
         require 'switchman/rails'
@@ -101,7 +102,7 @@ module Switchman
         self.default_shard = ::Rails.env.to_sym
         self.default_role = :primary
 
-        include ActiveRecord::Base
+        prepend ActiveRecord::Base
         include ActiveRecord::AttributeMethods
         include ActiveRecord::Persistence
         singleton_class.prepend ActiveRecord::ModelSchema::ClassMethods
@@ -150,9 +151,12 @@ module Switchman
         ::ActiveRecord::Relation.include(CallSuper)
 
         ::ActiveRecord::PredicateBuilder::AssociationQueryValue.prepend(ActiveRecord::PredicateBuilder::AssociationQueryValue)
+        ::ActiveRecord::PredicateBuilder::PolymorphicArrayValue.prepend(ActiveRecord::PredicateBuilder::AssociationQueryValue)
 
         ::ActiveRecord::Tasks::DatabaseTasks.singleton_class.prepend(ActiveRecord::Tasks::DatabaseTasks)
 
+        ::ActiveRecord::TestFixtures.prepend(ActiveRecord::TestFixtures)
+
         ::ActiveRecord::TypeCaster::Map.include(ActiveRecord::TypeCaster::Map)
         ::ActiveRecord::TypeCaster::Connection.include(ActiveRecord::TypeCaster::Connection)
 

data/lib/switchman/r_spec_helper.rb

@@ -111,28 +111,12 @@ module Switchman
           Shard.default(reload: true)
           @shard1 = Shard.find(@shard1.id)
           @shard2 = Shard.find(@shard2.id)
-          shards = [@shard2]
-          shards << @shard1 unless @shard1.database_server == Shard.default.database_server
-          shards.each do |shard|
-            shard.activate do
-              ::ActiveRecord::Base.connection.begin_transaction joinable: false
-            end
-          end
         end
       end
 
       klass.after do
         next if @@sharding_failed
 
-        if use_transactional_tests
-          shards = [@shard2]
-          shards << @shard1 unless @shard1.database_server == Shard.default.database_server
-          shards.each do |shard|
-            shard.activate do
-              ::ActiveRecord::Base.connection.rollback_transaction if ::ActiveRecord::Base.connection.transaction_open?
-            end
-          end
-        end
         # clean up after specs
         DatabaseServer.all.each do |ds|
           if ds.fake? && ds != @shard2.database_server
@@ -143,7 +127,8 @@ module Switchman
       end
 
       klass.after(:all) do
-        Shard.connection.update("TRUNCATE #{Shard.quoted_table_name} CASCADE")
+        # Don't truncate because that can create some fun cross-connection lock contention
+        Shard.delete_all
         Switchman.cache.delete('default_shard')
         Shard.default(reload: true)
       end

data/lib/switchman/rails.rb

@@ -13,7 +13,7 @@ module Switchman
       end
 
       def cache
-        Switchman::Shard.current.database_server.cache_store
+        Switchman::Shard.current ? Switchman::Shard.current.database_server.cache_store : super
       end
     end
   end

data/lib/switchman/standard_error.rb

@@ -10,10 +10,10 @@ module Switchman
         return
       end
 
-      if defined?(Shard)
-        @active_shards = Shard.sharded_models.map do |klass|
-          [klass, Shard.current(klass)]
-        end.compact.to_h
+      begin
+        @active_shards = Shard.active_shards if defined?(Shard)
+      rescue ::ActiveRecord::ConnectionNotEstablished
+        # If we hit an error really early in boot, activerecord may not be initialized yet
       end
 
       super

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = '3.0.2'
+  VERSION = '3.0.9'
 end

data/lib/tasks/switchman.rake

@@ -46,7 +46,7 @@ module Switchman
     end
 
     def self.options
-      { parallel: ENV['PARALLEL'].to_i, max_procs: ENV['MAX_PARALLEL_PROCS'] }
+      { parallel: ENV['PARALLEL'].to_i }
     end
 
     # classes - an array or proc, to activate as the current shard during the
@@ -80,14 +80,14 @@ module Switchman
             puts "Exception from #{e.current_shard.id}: #{e.current_shard.description}" if options[:parallel] != 0
             raise
 
-            #::ActiveRecord::Base.configurations = old_configurations
+            # ::ActiveRecord::Base.configurations = old_configurations
           end
         end
       end
     end
 
     %w[db:migrate db:migrate:up db:migrate:down db:rollback].each do |task_name|
-      shardify_task(task_name, classes: -> { Shard.sharded_models })
+      shardify_task(task_name)
     end
 
     def self.shard_scope(scope, raw_shard_ids)
@@ -201,14 +201,18 @@ module Switchman
     module PostgreSQLDatabaseTasks
       def structure_dump(filename, extra_flags = nil)
         set_psql_env
-        args = ['-s', '-x', '-O', '-f', filename]
+        args = ['--schema-only', '--no-privileges', '--no-owner', '--file', filename]
         args.concat(Array(extra_flags)) if extra_flags
         shard = Shard.current.name
         serialized_search_path = shard
         args << "--schema=#{Shellwords.escape(shard)}"
 
-        args << configuration['database']
+        ignore_tables = ::ActiveRecord::SchemaDumper.ignore_tables
+        args += ignore_tables.flat_map { |table| ['-T', table] } if ignore_tables.any?
+
+        args << db_config.database
         run_cmd('pg_dump', args, 'dumping')
+        remove_sql_header_comments(filename)
         File.open(filename, 'a') { |f| f << "SET search_path TO #{serialized_search_path};\n\n" }
       end
     end

metadata

@@ -1,16 +1,16 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.0.2
+  version: 3.0.9
 platform: ruby
 authors:
 - Cody Cutrer
 - James Williams
 - Jacob Fugal
-autorequire:
+autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-06-09 00:00:00.000000000 Z
+date: 2022-03-09 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -18,7 +18,7 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: 6.1.4
     - - "<"
       - !ruby/object:Gem::Version
         version: '6.2'
@@ -28,7 +28,7 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: 6.1.4
     - - "<"
       - !ruby/object:Gem::Version
         version: '6.2'
@@ -86,14 +86,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.1'
+        version: 2.3.0
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.1'
+        version: 2.3.0
 - !ruby/object:Gem::Dependency
   name: byebug
   requirement: !ruby/object:Gem::Requirement
@@ -275,6 +275,7 @@ files:
 - lib/switchman/active_record/statement_cache.rb
 - lib/switchman/active_record/table_definition.rb
 - lib/switchman/active_record/tasks/database_tasks.rb
+- lib/switchman/active_record/test_fixtures.rb
 - lib/switchman/active_record/type_caster.rb
 - lib/switchman/active_support/cache.rb
 - lib/switchman/arel.rb
@@ -297,8 +298,9 @@ files:
 homepage: http://www.instructure.com/
 licenses:
 - MIT
-metadata: {}
-post_install_message:
+metadata:
+  rubygems_mfa_required: 'true'
+post_install_message: 
 rdoc_options: []
 require_paths:
 - lib
@@ -313,8 +315,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.2.15
-signing_key:
+rubygems_version: 3.1.4
+signing_key: 
 specification_version: 4
 summary: Rails sharding magic
 test_files: []