switchman 1.5.21 → 2.1.5

data/db/migrate/20130328212039_create_switchman_shards.rb

@@ -1,4 +1,6 @@
-class CreateSwitchmanShards < ActiveRecord::Migration
+# frozen_string_literal: true
+
+class CreateSwitchmanShards < ActiveRecord::Migration[4.2]
   def change
     create_table :switchman_shards do |t|
       t.string :name

data/db/migrate/20130328224244_create_default_shard.rb

@@ -1,9 +1,11 @@
-class CreateDefaultShard < ActiveRecord::Migration
+# frozen_string_literal: true
+
+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(true)
+      Switchman::Shard.default(reload: true)
     end
   end
 end

data/db/migrate/20161206323434_add_back_default_string_limits_switchman.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+class AddBackDefaultStringLimitsSwitchman < ActiveRecord::Migration[4.2]
+  def up
+    add_string_limit_if_missing :switchman_shards, :name
+    add_string_limit_if_missing :switchman_shards, :database_server_id
+  end
+
+  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

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+class AddDefaultShardIndex < ActiveRecord::Migration[4.2]
+  def change
+    Switchman::Shard.where(default: nil).update_all(default: false)
+    change_column_default :switchman_shards, :default, false
+    change_column_null :switchman_shards, :default, false
+    options = if connection.adapter_name == 'PostgreSQL'
+                { unique: true, where: "\"default\"" }
+              else
+                {}
+              end
+    add_index :switchman_shards, :default, **options
+  end
+end

data/db/migrate/20180828192111_add_timestamps_to_shards.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+class AddTimestampsToShards < ActiveRecord::Migration[4.2]
+  def change
+    add_timestamps :switchman_shards, null: 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
+  end
+end

data/db/migrate/20190114212900_add_unique_name_indexes.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+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'
+  end
+end

data/lib/switchman/action_controller/caching.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActionController
     module Caching

data/lib/switchman/active_record/abstract_adapter.rb

@@ -1,11 +1,13 @@
+# frozen_string_literal: true
+
 require 'switchman/sharded_instrumenter'
 
 module Switchman
   module ActiveRecord
     module AbstractAdapter
       module ForeignKeyCheck
-        def add_column(table, name, type, options = {})
-          Engine.foreign_key_check(name, type, options)
+        def add_column(table, name, type, limit: nil, **)
+          Engine.foreign_key_check(name, type, limit: limit)
           super
         end
       end
@@ -27,8 +29,8 @@ module Switchman
         quote_table_name(name)
       end
 
-      def use_qualified_names?
-        false
+      def schema_migration
+        ::ActiveRecord::SchemaMigration
       end
 
       protected
@@ -38,6 +40,14 @@ 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

@@ -1,13 +1,10 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Association
       def shard
-        if @reflection.options[:polymorphic] || @reflection.klass.shard_category == @owner.class.shard_category
-          # polymorphic associations assume the same shard as the owning item
-          @owner.shard
-        else
-          Shard.default
-        end
+        reflection.shard(owner)
       end
 
       def build_record(*args)
@@ -18,12 +15,6 @@ module Switchman
         self.shard.activate { super }
       end
 
-      def association_scope
-        if klass
-          shard.activate(klass.shard_category) { super }
-        end
-      end
-
       def scope
         shard_value = @reflection.options[:multishard] ? @owner : self.shard
         @owner.shard.activate { super.shard(shard_value, :association) }
@@ -41,9 +32,11 @@ module Switchman
     end
 
     module CollectionAssociation
-      def get_records
-        shards = reflection.options[:multishard] && owner.respond_to?(:associated_shards) ? owner.associated_shards : [owner.shard]
-        Shard.with_each_shard(shards, [klass.shard_category]) do
+      def find_target
+        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
           super
         end
       end
@@ -69,30 +62,49 @@ module Switchman
       end
     end
 
-    if ::Rails.version >= '5'
-      module Extension
-        def self.build(_model, _reflection)
-        end
-
-        def self.valid_options
-          [:multishard]
-        end
+    module Extension
+      def self.build(_model, _reflection)
       end
 
-      ::ActiveRecord::Associations::Builder::Association.extensions << Extension
-    else
-      module Builder
-        module CollectionAssociation
-          def valid_options
-            super + [:multishard]
-          end
-        end
+      def self.valid_options
+        [:multishard]
       end
     end
 
+    ::ActiveRecord::Associations::Builder::Association.extensions << Extension
+
     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
+          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
 
           if klass.nil? || owners_map.empty?
@@ -140,12 +152,13 @@ module Switchman
             records.flatten!
           end
 
-          if ::Rails.version >= '4.1'
-            @preloaded_records = records
-          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
 
           # Each record may have multiple owners, and vice-versa
-          records_by_owner = owners.each_with_object({}) do |owner,h|
+          @associated_records_by_owner = owners.each_with_object({}) do |owner,h|
             h[owner] = []
           end
           records.each do |record|
@@ -153,10 +166,11 @@ module Switchman
             owner_key = Shard.global_id_for(owner_key, record.shard) if owner_key && record.class.sharded_column?(association_key_name)
 
             owners_map[owner_key.to_s].each do |owner|
-              records_by_owner[owner] << record
+              owner.association(reflection.name).set_inverse_instance(record)
+              @associated_records_by_owner[owner] << record
             end
           end
-          records_by_owner
+          @associated_records_by_owner
         end
 
         def owners_by_key
@@ -174,6 +188,12 @@ module Switchman
     end
 
     module CollectionProxy
+      def initialize(*args)
+        super
+        self.shard_value = scope.shard_value
+        self.shard_source_value = :association
+      end
+
       def shard(*args)
         scope.shard(*args)
       end
@@ -185,6 +205,13 @@ module Switchman
           (record.has_attribute?(reflection.foreign_key) && record.send(reflection.foreign_key) != key) || # have to use send instead of [] because sharding
           record.attribute_changed?(reflection.foreign_key)
       end
+
+      def save_belongs_to_association(reflection)
+        # 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 }
+      end
     end
   end
 end

data/lib/switchman/active_record/attribute_methods.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module AttributeMethods
@@ -27,7 +29,7 @@ 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 }.try(: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
@@ -36,11 +38,15 @@ module Switchman
         def define_method_global_attribute(attr_name)
           if sharded_column?(attr_name)
             generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
-              def __temp__
+              def __temp_global_attribute__
+                raw_value = original_#{attr_name}
+                return nil if raw_value.nil?
+                return raw_value if raw_value > Shard::IDS_PER_SHARD
+
                 Shard.global_id_for(original_#{attr_name}, shard)
               end
-              alias_method 'global_#{attr_name}', :__temp__
-              undef_method :__temp__
+              alias_method 'global_#{attr_name}', :__temp_global_attribute__
+              undef_method :__temp_global_attribute__
             RUBY
           else
             define_method_unsharded_column(attr_name, 'global')
@@ -50,23 +56,28 @@ module Switchman
         def define_method_local_attribute(attr_name)
           if sharded_column?(attr_name)
             generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
-              def __temp__
-                Shard.local_id_for(original_#{attr_name}).first
+              def __temp_local_attribute__
+                raw_value = original_#{attr_name}
+                return nil if raw_value.nil?
+                return raw_value % Shard::IDS_PER_SHARD
               end
-              alias_method 'local_#{attr_name}', :__temp__
-              undef_method :__temp__
+              alias_method 'local_#{attr_name}', :__temp_local_attribute__
+              undef_method :__temp_local_attribute__
             RUBY
           else
             define_method_unsharded_column(attr_name, 'local')
           end
         end
 
+        # see also Base#shard_category_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)
           if reflection
             if reflection.options[:polymorphic]
               # a polymorphic association has to be discovered at runtime. This code ends up being something like
-              # context_type.try(:constantize).try(:shard_category) || :default
-              "read_attribute(:#{reflection.foreign_type}).try(:constantize).try(:shard_category) || :default"
+              # context_type.&.constantize&.shard_category || :primary
+              "read_attribute(:#{reflection.foreign_type})&.constantize&.shard_category || :primary"
             else
               # otherwise we can just return a symbol for the statically known type of the association
               reflection.klass.shard_category.inspect
@@ -79,22 +90,43 @@ module Switchman
         def define_method_original_attribute(attr_name)
           if sharded_column?(attr_name)
             reflection = reflection_for_integer_attribute(attr_name)
-            generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
+            if attr_name == "id" && ::Rails.version >= '5.1.2'
+              return if self.method_defined?(:original_id)
+              owner = self
+            else
+              owner = generated_attribute_methods
+            end
+            owner.module_eval <<-RUBY, __FILE__, __LINE__ + 1
               # 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 __temp_relative_attribute__
+                raw_value = original_#{attr_name}
+                return nil if raw_value.nil?
+
+                abs_raw_value = raw_value.abs
+                current_shard = Shard.current(#{shard_category_code_for_reflection(reflection)})
+                same_shard = shard == current_shard
+                return raw_value if same_shard && abs_raw_value < Shard::IDS_PER_SHARD
+
+                value_shard_id = abs_raw_value / 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 % Shard::IDS_PER_SHARD if value_shard_id == current_shard.id
+                return raw_value if !same_shard && abs_raw_value > Shard::IDS_PER_SHARD
+                return shard.global_id_for(raw_value) if !same_shard && abs_raw_value < Shard::IDS_PER_SHARD
+
+                Shard.relative_id_for(raw_value, shard, current_shard)
               end
-              alias_method '#{attr_name}', :__temp__
-              undef_method :__temp__
+              alias_method '#{attr_name}', :__temp_relative_attribute__
+              undef_method :__temp_relative_attribute__
 
               alias_method 'original_#{attr_name}=', '#{attr_name}='
-              def __temp__(new_value)
+              def __temp_relative_attribute_assignment__(new_value)
                 self.original_#{attr_name} = Shard.relative_id_for(new_value, Shard.current(#{shard_category_code_for_reflection(reflection)}), shard)
               end
-              alias_method '#{attr_name}=', :__temp__
-              undef_method :__temp__
+              alias_method '#{attr_name}=', :__temp_relative_attribute_assignment__
+              undef_method :__temp_relative_attribute_assignment__
             RUBY
           else
             define_method_unsharded_column(attr_name, 'global')
@@ -104,11 +136,11 @@ module Switchman
         def define_method_unsharded_column(attr_name, prefix)
           return if columns_hash["#{prefix}_#{attr_name}"]
           generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
-              def __temp__
+              def __temp_unsharded_attribute__
                 raise NoMethodError, "undefined method `#{prefix}_#{attr_name}'; are you missing an association?"
               end
-              alias_method '#{prefix}_#{attr_name}', :__temp__
-              undef_method :__temp__
+              alias_method '#{prefix}_#{attr_name}', :__temp_unsharded_attribute__
+              undef_method :__temp_unsharded_attribute__
           RUBY
         end
       end
@@ -117,7 +149,7 @@ module Switchman
         klass.extend(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

data/lib/switchman/active_record/base.rb

@@ -1,11 +1,18 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Base
       module ClassMethods
         delegate :shard, to: :all
 
+        def find_ids_in_ranges(opts={}, &block)
+          opts.reverse_merge!(:loose => true)
+          all.find_ids_in_ranges(opts, &block)
+        end
+
         def shard_category
-          @shard_category || (self.superclass < ::ActiveRecord::Base && self.superclass.shard_category) || :default
+          connection_specification_name.to_sym
         end
 
         def shard_category=(category)
@@ -14,34 +21,32 @@ module Switchman
             categories[shard_category].delete(self)
             categories.delete(shard_category) if categories[shard_category].empty?
           end
-          connection_handler.uninitialize_ar(self)
           categories[category] ||= []
           categories[category] << self
-          @shard_category = category
-          connection_handler.initialize_categories(superclass)
+          self.connection_specification_name = category.to_s
         end
 
         def integral_id?
           if @integral_id == nil
-            @integral_id = columns_hash[primary_key].try(:type) == :integer
+            @integral_id = columns_hash[primary_key]&.type == :integer
           end
           @integral_id
         end
 
-        def transaction(*args)
+        def transaction(**)
           if self != ::ActiveRecord::Base && current_scope
             current_scope.activate do
               db = Shard.current(shard_category).database_server
-              if ::Shackles.environment != db.shackles_environment
-                db.unshackle { super }
+              if ::GuardRail.environment != db.guard_rail_environment
+                db.unguard { super }
               else
                 super
               end
             end
           else
             db = Shard.current(shard_category).database_server
-            if ::Shackles.environment != db.shackles_environment
-              db.unshackle { super }
+            if ::GuardRail.environment != db.guard_rail_environment
+              db.unguard { super }
             else
               super
             end
@@ -65,21 +70,29 @@ module Switchman
             result
           end
         end
-      end
 
-      def self.included(klass)
-        klass.extend(ClassMethods)
-        klass.set_callback(:initialize, :before) do
-          unless @shard
-            if self.class.sharded_primary_key?
-              @shard = Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.shard_category))
-            else
-              @shard = Shard.current(self.class.shard_category)
+        def clear_query_caches_for_current_thread
+          ::ActiveRecord::Base.connection_handlers.each_value do |handler|
+            handler.connection_pool_list.each do |pool|
+              pool.connection(switch_shard: false).clear_query_cache if pool.active_connection?
             end
           end
         end
       end
 
+      def self.prepended(klass)
+        klass.singleton_class.prepend(ClassMethods)
+      end
+
+      def _run_initialize_callbacks
+        @shard ||= if self.class.sharded_primary_key?
+          Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.shard_category))
+        else
+          Shard.current(self.class.shard_category)
+        end
+        super
+      end
+
       def shard
         @shard || Shard.current(self.class.shard_category) || Shard.default
       end
@@ -94,22 +107,18 @@ module Switchman
         end
       end
 
-      def scope_class
-        self.class.base_class
-      end
-
-      def save(*args)
+      def save(*, **)
         @shard_set_in_stone = true
-        scope_class.shard(shard, :implicit).scoping { super }
+        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
       end
 
-      def save!(*args)
+      def save!(*, **)
         @shard_set_in_stone = true
-        scope_class.shard(shard, :implicit).scoping { super }
+        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
       end
 
       def destroy
-        scope_class.shard(shard, :implicit).scoping { super }
+        self.class.shard(shard, :implicit).scoping { super }
       end
 
       def clone
@@ -121,14 +130,14 @@ module Switchman
         result
       end
 
-      def transaction(options={}, &block)
+      def transaction(**kwargs, &block)
         shard.activate(self.class.shard_category) do
-          self.class.transaction(options, &block)
+          self.class.transaction(**kwargs, &block)
         end
       end
 
       def hash
-        self.class.sharded_primary_key? ? Shard.global_id_for(id).hash : super
+        self.class.sharded_primary_key? ? self.class.hash ^ Shard.global_id_for(id).hash : super
       end
 
       def to_param
@@ -141,6 +150,42 @@ module Switchman
         @shard_set_in_stone = false
         copy
       end
+
+      def quoted_id
+        return super unless self.class.sharded_primary_key?
+        # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
+        self.class.connection.quote(id)
+      end
+
+      def update_columns(*)
+        db = shard.database_server
+        if ::GuardRail.environment != db.guard_rail_environment
+          return db.unguard { super }
+        else
+          super
+        end
+      end
+
+      protected
+
+      # see also AttributeMethods#shard_category_code_for_reflection
+      def shard_category_for_reflection(reflection)
+        if reflection
+          if reflection.options[:polymorphic]
+            begin
+              read_attribute(reflection.foreign_type)&.constantize&.shard_category || :primary
+            rescue NameError
+              # in case someone is abusing foreign_type to not point to an actual class
+              :primary
+            end
+          else
+            # otherwise we can just return a symbol for the statically known type of the association
+            reflection.klass.shard_category
+          end
+        else
+          shard_category
+        end
+      end
     end
   end
 end

data/lib/switchman/active_record/batches.rb

@@ -1,8 +1,10 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Batches
       def batch_order
-        "#{connection.quote_local_table_name(table_name)}.#{quoted_primary_key} ASC"
+        ::Arel.sql("#{connection.quote_local_table_name(table_name)}.#{quoted_primary_key} ASC")
       end
     end
   end