switchman 1.6.1 → 2.0.9

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 record.class.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
@@ -61,12 +63,15 @@ module Switchman
           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,7 +84,13 @@ 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
@@ -117,7 +128,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,6 +70,14 @@ module Switchman
             result
           end
         end
+
+        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.included(klass)
@@ -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.current(self.class.shard_category).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

data/lib/switchman/active_record/calculations.rb

@@ -1,16 +1,15 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Calculations
-      CALL_SUPER = Object.new.freeze
-      private_constant :CALL_SUPER
 
       def pluck(*column_names)
-        return super(*column_names[1..-1]) if column_names.first.equal?(CALL_SUPER)
         target_shard = Shard.current(klass.shard_category)
         shard_count = 0
         result = self.activate do |relation, shard|
           shard_count += 1
-          results = relation.pluck(CALL_SUPER, *column_names)
+          results = relation.call_super(:pluck, Calculations, *column_names)
           if column_names.length > 1
             column_names.each_with_index do |column_name, idx|
               if klass.sharded_column?(column_name)
@@ -30,13 +29,12 @@ module Switchman
         result
       end
 
-      def execute_simple_calculation(operation, column_name, distinct, super_method: false)
-        return super(operation, column_name, distinct) if super_method
+      def execute_simple_calculation(operation, column_name, distinct)
         operation = operation.to_s.downcase
         if operation == "average"
           result = calculate_simple_average(column_name, distinct)
         else
-          result = self.activate{ |relation| relation.send(:execute_simple_calculation, operation, column_name, distinct, super_method: true) }
+          result = self.activate{ |relation| relation.call_super(:execute_simple_calculation, Calculations, operation, column_name, distinct) }
           if result.is_a?(Array)
             case operation
             when "count", "sum"
@@ -61,12 +59,12 @@ module Switchman
         initial_results = relation.activate{ |rel| klass.connection.select_all(rel) }
         if initial_results.is_a?(Array)
           initial_results.each do |r|
-            r["average"] = type_cast_calculated_value(r["average"], column_or_type_for(column_name), "average")
-            r["count"] = type_cast_calculated_value(r["count"], column_or_type_for(column_name), "count")
+            r["average"] = type_cast_calculated_value(r["average"], type_for(column_name), "average")
+            r["count"] = type_cast_calculated_value(r["count"], type_for(column_name), "count")
           end
           result = initial_results.map{|r| r["average"] * r["count"]}.sum / initial_results.map{|r| r["count"]}.sum
         else
-          result = type_cast_calculated_value(initial_results.first["average"], column_or_type_for(column_name), "average")
+          result = type_cast_calculated_value(initial_results.first["average"], type_for(column_name), "average")
         end
         result
       end
@@ -76,7 +74,7 @@ module Switchman
         opts = grouped_calculation_options(operation.to_s.downcase, column_name, distinct)
 
         relation = build_grouped_calculation_relation(opts)
-        target_shard = Shard.current(:default)
+        target_shard = Shard.current(:primary)
 
         rows = relation.activate do |rel, shard|
           calculated_data = klass.connection.select_all(rel)
@@ -89,14 +87,14 @@ module Switchman
 
           calculated_data.map do |row|
             row[opts[:aggregate_alias]] = type_cast_calculated_value(
-                row[opts[:aggregate_alias]], column_or_type_for(opts[:column_name]), opts[:operation])
+                row[opts[:aggregate_alias]], type_for(opts[:column_name]), opts[:operation])
             row['count'] = row['count'].to_i if opts[:operation] == 'average'
 
-            opts[:group_columns].each do |aliaz, column_or_type, column_name|
+            opts[:group_columns].each do |aliaz, type, column_name|
               if opts[:associated] && (aliaz == opts[:group_aliases].first)
                 row[aliaz] = key_records[Shard.relative_id_for(row[aliaz], shard, target_shard)]
               elsif column_name && @klass.sharded_column?(column_name)
-                row[aliaz] = Shard.relative_id_for(type_cast_calculated_value(row[aliaz], column_or_type), shard, target_shard)
+                row[aliaz] = Shard.relative_id_for(type_cast_calculated_value(row[aliaz], type), shard, target_shard)
               end
             end
             row
@@ -112,10 +110,6 @@ module Switchman
         field.respond_to?(:name) ? field.name.to_s : field.to_s.split('.').last
       end
 
-      def column_or_type_for(field)
-        ::Rails.version < '4.2' ? column_for(field) : type_for(field)
-      end
-
       def grouped_calculation_options(operation, column_name, distinct)
         opts = {:operation => operation, :column_name => column_name, :distinct => distinct}
 
@@ -129,9 +123,10 @@ module Switchman
           group_fields = group_attrs
         end
 
-        group_aliases = group_fields.map { |field| column_alias_for(field) }
+        # to_s is because Rails 5 returns a string but Rails 6 returns a symbol.
+        group_aliases = group_fields.map { |field| column_alias_for(field.downcase.to_s).to_s }
         group_columns = group_aliases.zip(group_fields).map { |aliaz, field|
-          [aliaz, column_or_type_for(field), column_name_for(field)]
+          [aliaz, type_for(field), column_name_for(field)]
         }
         opts.merge!(:association => association, :associated => associated,
          :group_aliases => group_aliases, :group_columns => group_columns,
@@ -167,7 +162,7 @@ module Switchman
               'count', opts[:distinct]).as('count')
         end
 
-        haves = ::Rails.version >= "5" ? having_clause.send(:predicates) : having_values
+        haves = having_clause.send(:predicates)
         select_values += select_values unless haves.empty?
         select_values.concat opts[:group_fields].zip(opts[:group_aliases]).map { |field,aliaz|
           if field.respond_to?(:as)