switchman 3.1.3 → 3.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 54ae135ce532aa5578592027f6ce46cf3aedecdf0a6fd097ae3597752f436bae
-  data.tar.gz: 02d88f28e64e487c6841940a383afd92a56378ea7c2e4b1156252c0bcf0fa71f
+  metadata.gz: 20c8d947a29c73c09af7d1c880ef03aff2185c927d30954bb304b4592a7b4051
+  data.tar.gz: 5973474196cc2c8f2957932a65261998de24a6be3653cedeff374c1fe1fa925f
 SHA512:
-  metadata.gz: 383075584faafb39abc64be4f6e77a99c4ca3dfd5acc1c57bae885ba71c8d13ce2cb0f89d9332fd8387f5100886a5a0ffbb6f6d53e382b0e54412aea6b72cf06
-  data.tar.gz: 875421e2c515a671ed8503fd07b55d22e16f098ddde635b5530a3ab4801c9bebd06a5e4e75d7a547b1e84e5628288dac9aefc0c6f9a23cedad80b1f8389bb5de
+  metadata.gz: 10f59abb9c44a3b6e8ec3e87419862f5959ec8086117172d6fc044911c0efb3d96ec320fc04630e482200fbf99da4a7a8d7a291137bbf9ce019f1e7a54f5eeb0
+  data.tar.gz: 72fc286e4f143bf9d8344aa6e4a0405161e476b23eada4fd956f2657a1911c4e93acad649e76cb17e444821060085110d12c2cfe0a5f115bf9b7ddc598fb17ea

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}"

data/lib/switchman/active_record/statement_cache.rb

@@ -25,7 +25,7 @@ module Switchman
       # we can make some assumptions about the shard source
       # (e.g. infer from the primary key or use the current shard)
 
-      def execute(*args)
+      def execute(*args, &block)
         params, connection = args
         klass = @klass
         target_shard = nil
@@ -40,7 +40,7 @@ module Switchman
 
         target_shard.activate(klass.connection_class_for_self) do
           sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)
-          klass.find_by_sql(sql, bind_values)
+          klass.find_by_sql(sql, bind_values, &block)
         end
       end
 
@@ -66,7 +66,11 @@ module Switchman
 
         def primary_value_index
           primary_ba_index = @bound_attributes.index do |ba|
-            ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.primary
+            if ba.value.is_a?(::ActiveRecord::StatementCache::Substitute)
+              ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.primary
+            else
+              false
+            end
           end
           @indexes.index(primary_ba_index) if primary_ba_index
         end

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.3'
+  VERSION = '3.2.1'
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.1.3
+  version: 3.2.1
 platform: ruby
 authors:
 - Cody Cutrer
@@ -10,7 +10,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2022-08-03 00:00:00.000000000 Z
+date: 2022-11-22 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord