switchman 2.1.0 → 3.0.6

data/app/models/switchman/unsharded_record.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module Switchman
+  class UnshardedRecord < ::ActiveRecord::Base
+    sharded_model
+  end
+end

data/db/migrate/20130328212039_create_switchman_shards.rb

@@ -5,7 +5,7 @@ class CreateSwitchmanShards < ActiveRecord::Migration[4.2]
     create_table :switchman_shards do |t|
       t.string :name
       t.string :database_server_id
-      t.boolean :default, :default => false, :null => false
+      t.boolean :default, default: false, null: false
     end
   end
 end

data/db/migrate/20130328224244_create_default_shard.rb

@@ -2,10 +2,10 @@
 
 class CreateDefaultShard < ActiveRecord::Migration[4.2]
   def up
-    unless Switchman::Shard.default.is_a?(Switchman::Shard)
-      Switchman::Shard.reset_column_information
-      Switchman::Shard.create!(:default => true)
-      Switchman::Shard.default(reload: true)
-    end
+    return if Switchman::Shard.default.is_a?(Switchman::Shard)
+
+    Switchman::Shard.reset_column_information
+    Switchman::Shard.create!(default: true)
+    Switchman::Shard.default(reload: true)
   end
 end

data/db/migrate/20161206323434_add_back_default_string_limits_switchman.rb

@@ -8,6 +8,7 @@ class AddBackDefaultStringLimitsSwitchman < ActiveRecord::Migration[4.2]
 
   def add_string_limit_if_missing(table, column)
     return if column_exists?(table, column, :string, limit: 255)
+
     change_column table, column, :string, limit: 255
   end
 end

data/db/migrate/20180828183945_add_default_shard_index.rb

@@ -2,11 +2,11 @@
 
 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'
-                { unique: true, where: "\"default\"" }
+                { unique: true, where: '"default"' }
               else
                 {}
               end

data/db/migrate/20180828192111_add_timestamps_to_shards.rb

@@ -1,17 +1,19 @@
 # frozen_string_literal: true
 
 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
     change_column_null :switchman_shards, :created_at, false
 
-    if Switchman::Shard.current.default?
-      Switchman::Shard.connection.schema_cache.clear!
-      Switchman::Shard.reset_column_information
-      Switchman::Shard.columns
-    end
+    return unless Switchman::Shard.current.default?
+
+    Switchman::Shard.connection.schema_cache.clear!
+    Switchman::Shard.reset_column_information
+    Switchman::Shard.columns
   end
 end

data/db/migrate/20190114212900_add_unique_name_indexes.rb

@@ -2,8 +2,10 @@
 
 class AddUniqueNameIndexes < ActiveRecord::Migration[4.2]
   def change
-    add_index :switchman_shards, [:database_server_id, :name], unique: true
-    add_index :switchman_shards, :database_server_id, unique: true, where: "name IS NULL", name: 'index_switchman_shards_unique_primary_shard'
-    add_index :switchman_shards, "(true)", unique: true, where: "database_server_id IS NULL AND name IS NULL", name: 'index_switchman_shards_unique_primary_db_and_shard'
+    add_index :switchman_shards, %i[database_server_id name], unique: true
+    add_index :switchman_shards, :database_server_id, unique: true, where: 'name IS NULL',
+                                                      name: 'index_switchman_shards_unique_primary_shard'
+    add_index :switchman_shards, '(true)', unique: true, where: 'database_server_id IS NULL AND name IS NULL',
+                                           name: 'index_switchman_shards_unique_primary_db_and_shard'
   end
 end

data/lib/switchman/action_controller/caching.rb

@@ -13,7 +13,7 @@ module Switchman
         # disallow assigning to ActionController::Base.cache_store or
         # ActionController::Base#cache_store for the same reasons we disallow
         # assigning to Rails.cache
-        def cache_store=(cache)
+        def cache_store=(_cache)
           raise NoMethodError
         end
       end
@@ -21,7 +21,7 @@ module Switchman
       include ConfigMethods
 
       def self.included(base)
-        base.extend(ConfigMethods)
+        base.singleton_class.prepend(ConfigMethods)
       end
     end
   end

data/lib/switchman/active_record/abstract_adapter.rb

@@ -40,14 +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/association.rb

@@ -8,27 +8,17 @@ module Switchman
       end
 
       def build_record(*args)
-        self.shard.activate { super }
+        shard.activate { super }
       end
 
       def load_target
-        self.shard.activate { super }
+        shard.activate { super }
       end
 
       def scope
-        shard_value = @reflection.options[:multishard] ? @owner : self.shard
+        shard_value = @reflection.options[:multishard] ? @owner : shard
         @owner.shard.activate { super.shard(shard_value, :association) }
       end
-
-      def creation_attributes
-        attributes = super
-
-        # translate keys
-        if reflection.macro.in?([:has_one, :has_many]) && !options[:through]
-          attributes[reflection.foreign_key] = Shard.relative_id_for(owner[reflection.active_record_primary_key], owner.shard, self.shard)
-        end
-        attributes
-      end
     end
 
     module CollectionAssociation
@@ -36,15 +26,19 @@ module Switchman
         shards = reflection.options[:multishard] && owner.respond_to?(:associated_shards) ? owner.associated_shards : [shard]
         # activate both the owner and the target's shard category, so that Reflection#join_id_for,
         # when called for the owner, will be returned relative to shard the query will execute on
-        Shard.with_each_shard(shards, [klass.shard_category, owner.class.shard_category].uniq) do
+        Shard.with_each_shard(shards, [klass.connection_classes, owner.class.connection_classes].uniq) do
           super
         end
       end
+
+      def _create_record(*)
+        shard.activate { super }
+      end
     end
 
     module BelongsToAssociation
-      def replace_keys(record)
-        if record && record.class.sharded_column?(reflection.association_primary_key(record.class))
+      def replace_keys(record, force: false)
+        if record&.class&.sharded_column?(reflection.association_primary_key(record.class))
           foreign_id = record[reflection.association_primary_key(record.class)]
           owner[reflection.foreign_key] = Shard.relative_id_for(foreign_id, record.shard, owner.shard)
         else
@@ -54,7 +48,7 @@ module Switchman
 
       def shard
         if @owner.class.sharded_column?(@reflection.foreign_key) &&
-            foreign_id = @owner[@reflection.foreign_key]
+           (foreign_id = @owner[@reflection.foreign_key])
           Shard.shard_for(foreign_id, @owner.shard)
         else
           super
@@ -62,9 +56,22 @@ module Switchman
       end
     end
 
-    module Extension
-      def self.build(_model, _reflection)
+    module ForeignAssociation
+      # significant change:
+      #   * transpose the key to the correct shard
+      def set_owner_attributes(record) # rubocop:disable Naming/AccessorMethodName
+        return if options[:through]
+
+        key = owner._read_attribute(reflection.join_foreign_key)
+        key = Shard.relative_id_for(key, owner.shard, shard)
+        record._write_attribute(reflection.join_primary_key, key)
+
+        record._write_attribute(reflection.type, owner.class.polymorphic_name) if reflection.type
       end
+    end
+
+    module Extension
+      def self.build(_model, _reflection); end
 
       def self.valid_options
         [:multishard]
@@ -75,45 +82,35 @@ module Switchman
 
     module Preloader
       module Association
-        if ::Rails.version >= "5.2" and ::Rails.version < "6.0"
-          def run(preloader)
-            associated_records_by_owner.each do |owner, records|
-              associate_records_to_owner(owner, records)
-            end
+        # Copypasta from Activerecord but with added global_id_for goodness.
+        def records_for(ids)
+          scope.where(association_key_name => ids).load do |record|
+            global_key = if model.connection_classes == UnshardedRecord
+                           convert_key(record[association_key_name])
+                         else
+                           Shard.global_id_for(record[association_key_name], record.shard)
+                         end
+            owner = owners_by_key[convert_key(global_key)].first
+            association = owner.association(reflection.name)
+            association.set_inverse_instance(record)
           end
         end
 
-        if ::Rails.version >= "6.0"
-          # Copypasta from Activerecord but with added global_id_for goodness.
-          def records_for(ids)
-            scope.where(association_key_name => ids).load do |record|
-              global_key = if model.shard_category == :unsharded
-                             convert_key(record[association_key_name])
-                           else
-                             Shard.global_id_for(record[association_key_name], record.shard)
-                           end
-              owner = owners_by_key[global_key.to_s].first
-              association = owner.association(reflection.name)
-              association.set_inverse_instance(record)
-            end
-          end
-
-          def records_by_owner
-            associated_records_by_owner
-          end
-        end
-
-        def associated_records_by_owner(preloader = nil)
-          return @associated_records_by_owner if defined?(@associated_records_by_owner)
-          owners_map = owners_by_key
+        # significant changes:
+        #  * partition_by_shard the records_for call
+        #  * re-globalize the fetched owner id before looking up in the map
+        def load_records
+          # owners can be duplicated when a relation has a collection association join
+          # #compare_by_identity makes such owners different hash keys
+          @records_by_owner = {}.compare_by_identity
 
-          if klass.nil? || owners_map.empty?
-            records = []
+          if owner_keys.empty?
+            raw_records = []
           else
             # determine the shard to search for each owner
             if reflection.macro == :belongs_to
               # for belongs_to, it's the shard of the foreign_key
-              partition_proc = ->(owner) do
+              partition_proc = lambda do |owner|
                 if owner.class.sharded_column?(owner_key_name)
                   Shard.shard_for(owner[owner_key_name], owner.shard)
                 else
@@ -123,64 +120,56 @@ module Switchman
             elsif !reflection.options[:multishard]
               # for non-multishard associations, it's *just* the owner's shard
               partition_proc = ->(owner) { owner.shard }
-            else
-              # for multishard associations, it's the owner object itself
-              # (all associated shards)
-
-              # this is the default behavior of partition_by_shard, so just let it be nil
-              # to avoid the proc call
-              # partition_proc = ->(owner) { owner }
             end
 
-            records = Shard.partition_by_shard(owners, partition_proc) do |partitioned_owners|
-              # Some databases impose a limit on the number of ids in a list (in Oracle it's 1000)
-              # Make several smaller queries if necessary or make one query if the adapter supports it
-              sliced_owners = partitioned_owners.each_slice(model.connection.in_clause_length || partitioned_owners.size)
-              sliced_owners.map do |slice|
-                relative_owner_keys = slice.map do |owner|
-                  key = owner[owner_key_name]
-                  if key && owner.class.sharded_column?(owner_key_name)
-                    key = Shard.relative_id_for(key, owner.shard, Shard.current(owner.class.shard_category))
-                  end
-                  key && key.to_s
+            raw_records = Shard.partition_by_shard(owners, partition_proc) do |partitioned_owners|
+              relative_owner_keys = partitioned_owners.map do |owner|
+                key = owner[owner_key_name]
+                if key && owner.class.sharded_column?(owner_key_name)
+                  key = Shard.relative_id_for(key, owner.shard,
+                                              Shard.current(klass.connection_classes))
                 end
-                relative_owner_keys.compact!
-                relative_owner_keys.uniq!
-                records_for(relative_owner_keys)
+                convert_key(key)
               end
+              relative_owner_keys.compact!
+              relative_owner_keys.uniq!
+              records_for(relative_owner_keys)
             end
-            records.flatten!
           end
 
-          # This ivar may look unused, but remember this is an extension of
-          # rails' AR::Associations::Preloader::Association class. It gets used
-          # by that class (and its subclasses).
-          @preloaded_records = records
+          @preloaded_records = raw_records.select do |record|
+            assignments = false
 
-          # Each record may have multiple owners, and vice-versa
-          @associated_records_by_owner = owners.each_with_object({}) do |owner,h|
-            h[owner] = []
-          end
-          records.each do |record|
             owner_key = record[association_key_name]
-            owner_key = Shard.global_id_for(owner_key, record.shard) if owner_key && record.class.sharded_column?(association_key_name)
+            if owner_key && record.class.sharded_column?(association_key_name)
+              owner_key = Shard.global_id_for(owner_key,
+                                              record.shard)
+            end
 
-            owners_map[owner_key.to_s].each do |owner|
-              owner.association(reflection.name).set_inverse_instance(record)
-              @associated_records_by_owner[owner] << record
+            owners_by_key[convert_key(owner_key)].each do |owner|
+              entries = (@records_by_owner[owner] ||= [])
+
+              if reflection.collection? || entries.empty?
+                entries << record
+                assignments = true
+              end
             end
+
+            assignments
           end
-          @associated_records_by_owner
         end
 
+        # significant change: globalize keys on sharded columns
         def owners_by_key
-          @owners_by_key ||= owners.group_by do |owner|
+          @owners_by_key ||= owners.each_with_object({}) do |owner, result|
             key = owner[owner_key_name]
             key = Shard.global_id_for(key, owner.shard) if key && owner.class.sharded_column?(owner_key_name)
-            key && key.to_s
+            key = convert_key(key)
+            (result[key] ||= []) << owner if key
           end
         end
 
+        # significant change: don't cache scope (since it could be for different shards)
         def scope
           build_scope
         end
@@ -210,7 +199,7 @@ module Switchman
         # this seems counter-intuitive, but the autosave code will assign to attribute bypassing switchman,
         # after reading the id attribute _without_ bypassing switchman. So we need Shard.current for the
         # category of the associated record to match Shard.current for the category of self
-        shard.activate(shard_category_for_reflection(reflection)) { super }
+        shard.activate(connection_classes_for_reflection(reflection)) { super }
       end
     end
   end

data/lib/switchman/active_record/attribute_methods.rb

@@ -4,22 +4,23 @@ module Switchman
   module ActiveRecord
     module AttributeMethods
       module ClassMethods
-
         def sharded_primary_key?
-          self != Shard && shard_category != :unsharded && integral_id?
+          !(self <= UnshardedRecord) && integral_id?
         end
 
         def sharded_foreign_key?(column_name)
           reflection = reflection_for_integer_attribute(column_name.to_s)
           return false unless reflection
+
           reflection.options[:polymorphic] || reflection.klass.sharded_primary_key?
         end
 
         def sharded_column?(column_name)
           column_name = column_name.to_s
           @sharded_column_values ||= {}
-          unless @sharded_column_values.has_key?(column_name)
-            @sharded_column_values[column_name] = (column_name == primary_key && sharded_primary_key?) || sharded_foreign_key?(column_name)
+          unless @sharded_column_values.key?(column_name)
+            @sharded_column_values[column_name] =
+              (column_name == primary_key && sharded_primary_key?) || sharded_foreign_key?(column_name)
           end
           @sharded_column_values[column_name]
         end
@@ -29,112 +30,186 @@ module Switchman
         def reflection_for_integer_attribute(attr_name)
           attr_name = attr_name.to_s
           columns_hash[attr_name] && columns_hash[attr_name].type == :integer &&
-              reflections.find { |_, r| r.belongs_to? && r.foreign_key.to_s == attr_name }&.last
+            reflections.find { |_, r| r.belongs_to? && r.foreign_key.to_s == attr_name }&.last
         rescue ::ActiveRecord::StatementInvalid
           # this is for when models are referenced in initializers before migrations have been run
-          raise if connection.open_transactions > 0
+          raise if connection.open_transactions.positive?
         end
 
-        def define_method_global_attribute(attr_name)
+        def define_method_global_attribute(attr_name, owner:)
           if sharded_column?(attr_name)
-            generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
-              def __temp__
-                Shard.global_id_for(original_#{attr_name}, shard)
+            owner << <<-RUBY
+              def global_#{attr_name}
+                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
-              alias_method 'global_#{attr_name}', :__temp__
-              undef_method :__temp__
             RUBY
           else
-            define_method_unsharded_column(attr_name, 'global')
+            define_method_unsharded_column(attr_name, 'global', owner)
           end
         end
 
-        def define_method_local_attribute(attr_name)
+        def define_method_local_attribute(attr_name, owner:)
           if sharded_column?(attr_name)
-            generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
-              def __temp__
-                Shard.local_id_for(original_#{attr_name}).first
+            owner << <<-RUBY
+              def local_#{attr_name}
+                raw_value = original_#{attr_name}
+                return nil if raw_value.nil?
+                return raw_value % ::Switchman::Shard::IDS_PER_SHARD
               end
-              alias_method 'local_#{attr_name}', :__temp__
-              undef_method :__temp__
             RUBY
           else
-            define_method_unsharded_column(attr_name, 'local')
+            define_method_unsharded_column(attr_name, 'local', owner)
           end
         end
 
-        # see also Base#shard_category_for_reflection
+        # see also Base#connection_classes_for_reflection
         # the difference being this will output static strings for the common cases, making them
         # more performant
-        def shard_category_code_for_reflection(reflection)
+        def connection_classes_code_for_reflection(reflection)
           if reflection
             if reflection.options[:polymorphic]
               # a polymorphic association has to be discovered at runtime. This code ends up being something like
-              # context_type.&.constantize&.shard_category || :primary
-              "read_attribute(:#{reflection.foreign_type})&.constantize&.shard_category || :primary"
+              # context_type.&.constantize&.connection_classes
+              "read_attribute(:#{reflection.foreign_type})&.constantize&.connection_classes"
             else
               # otherwise we can just return a symbol for the statically known type of the association
-              reflection.klass.shard_category.inspect
+              "::#{reflection.klass.connection_classes.name}"
             end
           else
-            shard_category.inspect
+            "::#{connection_classes.name}"
+          end
+        end
+
+        # just a dummy class with the proper interface that calls module_eval immediately
+        class CodeGenerator
+          def initialize(mod, line)
+            @module = mod
+            @line = line
+          end
+
+          def <<(string)
+            @module.module_eval(string, __FILE__, @line)
           end
         end
 
-        def define_method_original_attribute(attr_name)
+        def define_method_original_attribute(attr_name, owner:)
           if sharded_column?(attr_name)
             reflection = reflection_for_integer_attribute(attr_name)
-            if attr_name == "id" && ::Rails.version >= '5.1.2'
-              return if self.method_defined?(:original_id)
-              owner = self
-            else
-              owner = generated_attribute_methods
+            if attr_name == 'id'
+              return if method_defined?(:original_id)
+
+              owner = CodeGenerator.new(self, __LINE__ + 4)
             end
-            owner.module_eval <<-RUBY, __FILE__, __LINE__ + 1
-              # rename the original method to original_
+
+            owner << <<-RUBY
+              # rename the original method to original_*
               alias_method 'original_#{attr_name}', '#{attr_name}'
               # and replace with one that transposes the id
-              def __temp__
-                Shard.relative_id_for(original_#{attr_name}, shard, Shard.current(#{shard_category_code_for_reflection(reflection)}))
+              def #{attr_name}
+                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 '#{attr_name}', :__temp__
-              undef_method :__temp__
 
               alias_method 'original_#{attr_name}=', '#{attr_name}='
-              def __temp__(new_value)
-                self.original_#{attr_name} = Shard.relative_id_for(new_value, Shard.current(#{shard_category_code_for_reflection(reflection)}), shard)
+              def #{attr_name}=(new_value)
+                self.original_#{attr_name} = ::Switchman::Shard.relative_id_for(new_value, ::Switchman::Shard.current(#{connection_classes_code_for_reflection(reflection)}), shard)
               end
-              alias_method '#{attr_name}=', :__temp__
-              undef_method :__temp__
             RUBY
           else
-            define_method_unsharded_column(attr_name, 'global')
+            define_method_unsharded_column(attr_name, 'global', owner)
           end
         end
 
-        def define_method_unsharded_column(attr_name, prefix)
+        def define_method_unsharded_column(attr_name, prefix, owner)
           return if columns_hash["#{prefix}_#{attr_name}"]
-          generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
-              def __temp__
-                raise NoMethodError, "undefined method `#{prefix}_#{attr_name}'; are you missing an association?"
-              end
-              alias_method '#{prefix}_#{attr_name}', :__temp__
-              undef_method :__temp__
+
+          owner << <<-RUBY
+            def #{prefix}_#{attr_name}
+              raise NoMethodError, "undefined method `#{prefix}_#{attr_name}'; are you missing an association?"
+            end
           RUBY
         end
       end
 
       def self.included(klass)
-        klass.extend(ClassMethods)
-        klass.attribute_method_prefix "global_", "local_", "original_"
+        klass.singleton_class.include(ClassMethods)
+        klass.attribute_method_prefix 'global_', 'local_', 'original_'
       end
 
       # ensure that we're using the sharded attribute method
       # and not the silly one in AR::AttributeMethods::PrimaryKey
       def id
+        return super if is_a?(Shard)
+
         self.class.define_attribute_methods
         super
       end
+
+      # these are called if the specific methods haven't been defined yet
+      def attribute(attr_name)
+        return super unless self.class.sharded_column?(attr_name)
+
+        reflection = self.class.send(:reflection_for_integer_attribute, attr_name)
+        ::Switchman::Shard.relative_id_for(super, shard, ::Switchman::Shard.current(connection_classes_for_reflection(reflection)))
+      end
+
+      def attribute=(attr_name, new_value)
+        unless self.class.sharded_column?(attr_name)
+          super
+          return
+        end
+
+        reflection = self.class.send(:reflection_for_integer_attribute, attr_name)
+        super(::Switchman::Shard.relative_id_for(new_value, ::Switchman::Shard.current(connection_classes_for_reflection(reflection)), shard))
+      end
+
+      def global_attribute(attr_name)
+        if self.class.sharded_column?(attr_name)
+          ::Switchman::Shard.global_id_for(attribute(attr_name), shard)
+        else
+          attribute(attr_name)
+        end
+      end
+
+      def local_attribute(attr_name)
+        if self.class.sharded_column?(attr_name)
+          ::Switchman::Shard.local_id_for(attribute(attr_name), shard).first
+        else
+          attribute(attr_name)
+        end
+      end
+
+      private
+
+      def connection_classes_for_reflection(reflection)
+        if reflection
+          if reflection.options[:polymorphic]
+            read_attribute(reflection.foreign_type)&.constantize&.connection_classes
+          else
+            reflection.klass.connection_classes
+          end
+        else
+          self.class.connection_classes
+        end
+      end
     end
   end
 end