switchman 3.1.0 → 3.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 48b27124d5233adf585aed105d0b2522c76249bdd1e13fbcfd7e904241d3a624
-  data.tar.gz: fea1ccc0ad63f00633fb90bee1d7bd4606f6de0e18564250865c1b1dde61ef8d
+  metadata.gz: 39d8ce1bd79211b12278263fda3f0a88adc841051bd1e4f52e8d01ccd371978f
+  data.tar.gz: b126fe1b7b700fb424351d901c58b66f9719892deab08a2242de6cf6c5645b77
 SHA512:
-  metadata.gz: c05b445111dd0e4b16ecbfaea6ab84557fa2cbadc6102fa0af6fb43cff2f6bbb6606e825f76ed20f6b8a4a2a2966e31babadf66c5eb065a6c4d8a7eea6a8b92a
-  data.tar.gz: 0af9e15706142f90772d24e2187db93ec83331997bccbb25c1b89cec2dde1bd2414d2774fec3785c7bea1a1b06d3fa5ad49e15b716f0aa35794e2e5734c848a8
+  metadata.gz: 266b29135322d76e811feb72414a5c3a1621afe07e65175672092ddefd21296b73fa3fe3ce1b6ce0d34b2e5d1a071b186f0b859b44398ff65e3a19cc3e747830
+  data.tar.gz: e26ecce87b654d3dd8ae8caa3111b9cbbc89896ce00cca467ba4bd66527ddf76c2e0db1c4ebff30cd133dbcefc1580f4f6b2dc745ab795813cf85848893badca

data/lib/switchman/active_record/associations.rb

@@ -83,16 +83,62 @@ module Switchman
 
       module Preloader
         module Association
-          module LoaderQuery
-            def load_records_in_batch(loaders)
-              # While in theory loading multiple associations that end up being effectively the same would be nice
-              # it's not very switchman compatible, so just don't bother trying to use that logic
-              # raw_records = records_for(loaders)
+          # significant changes:
+          #  * associate shards with records
+          #  * look on all appropriate shards when loading records
+          module LoaderRecords
+            def populate_keys_to_load_and_already_loaded_records
+              @sharded_keys_to_load = {}
 
               loaders.each do |loader|
-                loader.load_records(nil)
-                loader.run
+                multishard = loader.send(:reflection).options[:multishard]
+                belongs_to = loader.send(:reflection).macro == :belongs_to
+                loader.owners_by_key.each do |key, owners|
+                  if (loaded_owner = owners.find { |owner| loader.loaded?(owner) })
+                    already_loaded_records_by_key[key] = loader.target_for(loaded_owner)
+                  else
+                    shard_set = @sharded_keys_to_load[key] ||= Set.new
+                    owner_key_name = loader.send(:owner_key_name)
+                    owners.each do |owner|
+                      if multishard && owner.respond_to?(:associated_shards)
+                        shard_set.merge(owner.associated_shards.map(&:id))
+                      elsif belongs_to && owner.class.sharded_column?(owner_key_name)
+                        shard_set.add(Shard.shard_for(owner[owner_key_name], owner.shard).id)
+                      elsif belongs_to
+                        shard_set.add(Shard.current.id)
+                      else
+                        shard_set.add(owner.shard.id)
+                      end
+                    end
+                  end
+                end
+              end
+
+              @sharded_keys_to_load.delete_if { |key, _shards| already_loaded_records_by_key.include?(key) }
+            end
+
+            def load_records
+              ret = []
+
+              shards_with_keys = @sharded_keys_to_load.each_with_object({}) do |(key, shards), h|
+                shards.each { |shard| (h[shard] ||= []) << key }
+              end
+
+              shards_with_keys.each do |shard, keys|
+                Shard.lookup(shard).activate do
+                  scope_was = loader_query.scope
+                  begin
+                    loader_query.instance_variable_set(:@scope, loader_query.scope.shard(Shard.current(loader_query.scope.model.connection_class_for_self)))
+                    ret += loader_query.load_records_for_keys(keys) do |record|
+                      loaders.each { |l| l.set_inverse(record) }
+                    end
+                  ensure
+                    loader_query.instance_variable_set(:@scope, scope_was)
+                  end
+                end
               end
+
+              ret
             end
           end
 
@@ -110,6 +156,26 @@ module Switchman
             end
           end
 
+          # Disabling to keep closer to rails original
+          # rubocop:disable Naming/AccessorMethodName, Style/GuardClause
+          # significant changes:
+          #  * globalize the key to lookup
+          def set_inverse(record)
+            global_key = if model.connection_class_for_self == UnshardedRecord
+                           convert_key(record[association_key_name])
+                         else
+                           Shard.global_id_for(record[association_key_name], record.shard)
+                         end
+
+            if (owners = owners_by_key[convert_key(global_key)])
+              # Processing only the first owner
+              # because the record is modified but not an owner
+              association = owners.first.association(reflection.name)
+              association.set_inverse_instance(record)
+            end
+          end
+          # rubocop:enable Naming/AccessorMethodName, Style/GuardClause
+
           # significant changes:
           #  * partition_by_shard the records_for call
           #  * re-globalize the fetched owner id before looking up in the map
@@ -120,7 +186,9 @@ module Switchman
             # #compare_by_identity makes such owners different hash keys
             @records_by_owner = {}.compare_by_identity
 
-            if ::Rails.version < '7.0' && owner_keys.empty?
+            if ::Rails.version >= '7.0'
+              raw_records ||= loader_query.records_for([self])
+            elsif owner_keys.empty?
               raw_records ||= []
             else
               # determine the shard to search for each owner
@@ -162,7 +230,7 @@ module Switchman
                                                 record.shard)
               end
 
-              owners_by_key[convert_key(owner_key)].each do |owner|
+              owners_by_key[convert_key(owner_key)]&.each do |owner|
                 entries = (@records_by_owner[owner] ||= [])
 
                 if reflection.collection? || entries.empty?

data/lib/switchman/active_record/attribute_methods.rb

@@ -99,7 +99,7 @@ module Switchman
             if reflection.options[:polymorphic]
               # a polymorphic association has to be discovered at runtime. This code ends up being something like
               # context_type.&.constantize&.connection_class_for_self
-              "read_attribute(:#{reflection.foreign_type})&.constantize&.connection_class_for_self"
+              "begin;read_attribute(:#{reflection.foreign_type})&.constantize&.connection_class_for_self;rescue NameError;::ActiveRecord::Base;end"
             else
               # otherwise we can just return a symbol for the statically known type of the association
               "::#{reflection.klass.connection_class_for_self.name}"
@@ -256,20 +256,6 @@ module Switchman
           attribute(attr_name)
         end
       end
-
-      private
-
-      def connection_class_for_self_for_reflection(reflection)
-        if reflection
-          if reflection.options[:polymorphic]
-            read_attribute(reflection.foreign_type)&.constantize&.connection_class_for_self
-          else
-            reflection.klass.connection_class_for_self
-          end
-        else
-          self.class.connection_class_for_self
-        end
-      end
     end
   end
 end

data/lib/switchman/active_record/base.rb

@@ -22,16 +22,20 @@ module Switchman
           @integral_id
         end
 
-        def transaction(**)
-          if self != ::ActiveRecord::Base && current_scope
-            current_scope.activate do
-              db = Shard.current(connection_class_for_self).database_server
-              db.unguard { super }
+        %w[transaction insert_all upsert_all].each do |method|
+          class_eval <<-RUBY, __FILE__, __LINE__ + 1
+            def #{method}(*, **)
+              if self != ::ActiveRecord::Base && current_scope
+                current_scope.activate do
+                  db = Shard.current(connection_class_for_self).database_server
+                  db.unguard { super }
+                end
+              else
+                db = Shard.current(connection_class_for_self).database_server
+                db.unguard { super }
+              end
             end
-          else
-            db = Shard.current(connection_class_for_self).database_server
-            db.unguard { super }
-          end
+          RUBY
         end
 
         def reset_column_information
@@ -138,9 +142,33 @@ module Switchman
                    else
                      Shard.current(self.class.connection_class_for_self)
                    end
+        readonly! if shadow_record? && !Switchman.config[:writable_shadow_records]
         super
       end
 
+      def shadow_record?
+        pkey = self[self.class.primary_key]
+        return false unless self.class.sharded_column?(self.class.primary_key) && pkey
+
+        pkey > Shard::IDS_PER_SHARD
+      end
+
+      def save_shadow_record(new_attrs: attributes, target_shard: Shard.current)
+        return if target_shard == shard
+
+        shadow_attrs = {}
+        new_attrs.each do |attr, value|
+          shadow_attrs[attr] = if self.class.sharded_column?(attr)
+                                 Shard.relative_id_for(value, shard, target_shard)
+                               else
+                                 value
+                               end
+        end
+        target_shard.activate do
+          self.class.upsert(shadow_attrs, unique_by: self.class.primary_key)
+        end
+      end
+
       def shard
         @shard || Shard.current(self.class.connection_class_for_self) || Shard.default
       end

data/lib/switchman/active_record/migration.rb

@@ -20,13 +20,15 @@ module Switchman
     end
 
     module Migrator
-      # 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
+      # significant change: use the shard name instead of the database name
+      # in the lock id. Especially if you're going through pgbouncer, the
+      # database name you're accessing may not be consistent
       def generate_migrator_advisory_lock_id
-        ::ActiveRecord::Migrator::MIGRATOR_SALT
+        db_name_hash = Zlib.crc32(Shard.current.name)
+        shard_name_hash = ::ActiveRecord::Migrator::MIGRATOR_SALT * db_name_hash
+        # Store in internalmetadata to allow other tools to be able to lock out migrations
+        ::ActiveRecord::InternalMetadata[:migrator_advisory_lock_id] = shard_name_hash
+        shard_name_hash
       end
 
       # significant change: strip out prefer_secondary from config

data/lib/switchman/engine.rb

@@ -50,7 +50,7 @@ module Switchman
         ::ActiveRecord::Associations::CollectionProxy.include(ActiveRecord::Associations::CollectionProxy)
 
         ::ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecord::Associations::Preloader::Association)
-        ::ActiveRecord::Associations::Preloader::Association::LoaderQuery.prepend(ActiveRecord::Associations::Preloader::Association::LoaderQuery) unless ::Rails.version < '7.0'
+        ::ActiveRecord::Associations::Preloader::Association::LoaderRecords.prepend(ActiveRecord::Associations::Preloader::Association::LoaderRecords) unless ::Rails.version < '7.0'
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::AbstractAdapter)
         ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ActiveRecord::ConnectionPool)
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::QueryCache)

data/lib/switchman/version.rb

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

metadata

@@ -1,16 +1,16 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.1.0
+  version: 3.2.0
 platform: ruby
 authors:
 - Cody Cutrer
 - James Williams
 - Jacob Fugal
-autorequire:
+autorequire: 
 bindir: bin
 cert_chain: []
-date: 2022-06-02 00:00:00.000000000 Z
+date: 2022-11-22 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -300,7 +300,7 @@ licenses:
 - MIT
 metadata:
   rubygems_mfa_required: 'true'
-post_install_message:
+post_install_message: 
 rdoc_options: []
 require_paths:
 - lib
@@ -316,7 +316,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
       version: '0'
 requirements: []
 rubygems_version: 3.1.6
-signing_key:
+signing_key: 
 specification_version: 4
 summary: Rails sharding magic
 test_files: []