switchman 1.13.3 → 2.2.0

data/db/migrate/20130328212039_create_switchman_shards.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class CreateSwitchmanShards < ActiveRecord::Migration[4.2]
   def change
     create_table :switchman_shards do |t|

data/db/migrate/20130328224244_create_default_shard.rb

@@ -1,9 +1,11 @@
+# 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

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class AddBackDefaultStringLimitsSwitchman < ActiveRecord::Migration[4.2]
   def up
     add_string_limit_if_missing :switchman_shards, :name

data/db/migrate/20180828183945_add_default_shard_index.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class AddDefaultShardIndex < ActiveRecord::Migration[4.2]
   def change
     Switchman::Shard.where(default: nil).update_all(default: false)
@@ -8,6 +10,6 @@ class AddDefaultShardIndex < ActiveRecord::Migration[4.2]
               else
                 {}
               end
-    add_index :switchman_shards, :default, options
+    add_index :switchman_shards, :default, **options
   end
 end

data/db/migrate/20180828192111_add_timestamps_to_shards.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class AddTimestampsToShards < ActiveRecord::Migration[4.2]
   def change
     add_timestamps :switchman_shards, null: true

data/db/migrate/20190114212900_add_unique_name_indexes.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class AddUniqueNameIndexes < ActiveRecord::Migration[4.2]
   def change
     add_index :switchman_shards, [:database_server_id, :name], unique: true

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

data/lib/switchman/active_record/association.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Association
@@ -30,17 +32,14 @@ module Switchman
     end
 
     module CollectionAssociation
-      method = ::Rails.version < '5.1' ? :get_records : :find_target
-      module_eval <<-RUBY, __FILE__, __LINE__ + 1
-        def #{method}
-          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
+      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
-      RUBY
+      end
     end
 
     module BelongsToAssociation
@@ -76,16 +75,36 @@ module Switchman
 
     module Preloader
       module Association
-        if ::Rails.version >= "5.2"
+        if ::Rails.version >= "5.2" and ::Rails.version < "6.0"
           def run(preloader)
-            # TODO - can move associated_records_by_owner into this after 5.1 is gonzo
-            associated_records_by_owner(preloader).each do |owner, records|
+            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?
@@ -133,10 +152,13 @@ module Switchman
             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
 
           # 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|
@@ -145,10 +167,10 @@ module Switchman
 
             owners_map[owner_key.to_s].each do |owner|
               owner.association(reflection.name).set_inverse_instance(record)
-              records_by_owner[owner] << record
+              @associated_records_by_owner[owner] << record
             end
           end
-          records_by_owner
+          @associated_records_by_owner
         end
 
         def owners_by_key
@@ -183,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
@@ -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,17 +56,22 @@ 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]
@@ -89,18 +100,33 @@ module Switchman
               # 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')
@@ -110,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

data/lib/switchman/active_record/base.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Base
@@ -31,20 +33,20 @@ module Switchman
           @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
@@ -68,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
@@ -97,14 +107,14 @@ module Switchman
         end
       end
 
-      def save(*args)
+      def save(*, **)
         @shard_set_in_stone = true
-        self.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
-        self.class.shard(shard, :implicit).scoping { super }
+        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
       end
 
       def destroy
@@ -120,9 +130,9 @@ 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
 
@@ -146,6 +156,36 @@ module Switchman
         # 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,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Batches

data/lib/switchman/active_record/calculations.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Calculations
@@ -49,7 +51,7 @@ module Switchman
 
       def calculate_simple_average(column_name, distinct)
         # See activerecord#execute_simple_calculation
-        relation = reorder(nil)
+        relation = except(:order)
         column = aggregate_column(column_name)
         relation.select_values = [operation_over_aggregate_column(column, "average", distinct).as("average"),
                                   operation_over_aggregate_column(column, "count", distinct).as("count")]
@@ -121,7 +123,8 @@ 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, type_for(field), column_name_for(field)]
         }

data/lib/switchman/active_record/connection_handler.rb

@@ -1,11 +1,11 @@
+# frozen_string_literal: true
+
 require 'switchman/connection_pool_proxy'
 
 module Switchman
   module ActiveRecord
     module ConnectionHandler
       def self.make_sharing_automagic(config, shard = Shard.current)
-        key = config[:adapter] == 'postgresql' ? :schema_search_path : :database
-
         # only load the shard name from the db if we have to
         if !config[:shard_name]
           # we may not be able to connect to this shard yet, cause it might be an empty database server
@@ -15,15 +15,13 @@ module Switchman
 
           config[:shard_name] ||= shard_name
         end
-
-        if !config[key] || config[key] == shard_name
-          # this may truncate the schema_search_path if it was not specified in database.yml
-          # but that's what our old behavior was anyway, so I guess it's okay
-          config[key] = '%{shard_name}'
-        end
       end
 
       def establish_connection(spec)
+        # Just skip establishing a sharded connection if sharding isn't loaded; we'll do it again later
+        # This only can happen when loading ActiveRecord::Base; after everything is loaded Shard will
+        # be defined and this will actually establish a connection
+        return unless defined?(Shard)
         pool = super
 
         # this is the first place that the adapter would have been required; but now we
@@ -39,19 +37,27 @@ module Switchman
           # to sharding will recurse onto itself trying to access column information
           Shard.default
 
+          config = pool.spec.config
           # automatically change config to allow for sharing connections with simple config
-          config = ::Rails.version < '5.1' ? spec.config : pool.spec.config
           ConnectionHandler.make_sharing_automagic(config)
           ConnectionHandler.make_sharing_automagic(Shard.default.database_server.config)
 
-          if ::Rails.version < '5.1'
-            ::ActiveRecord::Base.configurations[::Rails.env] = spec.instance_variable_get(:@config).stringify_keys
-          else
+          if ::Rails.version < '6.0'
             ::ActiveRecord::Base.configurations[::Rails.env] = config.stringify_keys
+          else
+            # Adopted from the deprecated code that currently lives in rails proper
+            remaining_configs = ::ActiveRecord::Base.configurations.configurations.reject { |db_config| db_config.env_name == ::Rails.env }
+            new_config = ::ActiveRecord::DatabaseConfigurations.new(::Rails.env => config.stringify_keys).configurations
+            new_configs = remaining_configs + new_config
+
+            ::ActiveRecord::Base.configurations = new_configs
           end
+        else
+          # this is probably wrong now
+          Shard.default.remove_instance_variable(:@name) if Shard.default.instance_variable_defined?(:@name)
         end
 
-        @shard_connection_pools ||= { [:master, Shard.default.database_server.shareable? ? ::Rails.env : Shard.default] => pool}
+        @shard_connection_pools ||= { [:primary, Shard.default.database_server.shareable? ? ::Rails.env : Shard.default] => pool}
 
         category = pool.spec.name.to_sym
         proxy = ConnectionPoolProxy.new(category,
@@ -60,13 +66,13 @@ module Switchman
         owner_to_pool[pool.spec.name] = proxy
 
         if first_time
-          if Shard.default.database_server.config[:prefer_slave]
-            Shard.default.database_server.shackle!
+          if Shard.default.database_server.config[:prefer_secondary]
+            Shard.default.database_server.guard!
           end
 
-          if Shard.default.is_a?(DefaultShard) && Shard.default.database_server.config[:slave]
-            Shard.default.database_server.shackle!
-            Shard.default(true)
+          if Shard.default.is_a?(DefaultShard) && Shard.default.database_server.config[:secondary]
+            Shard.default.database_server.guard!
+            Shard.default(reload: true)
           end
         end
 
@@ -106,8 +112,20 @@ module Switchman
       end
 
       def remove_connection(spec_name)
+        # also remove pools based on the same spec name that are for shard category purposes
+        # can't just use delete_if, because it's a Concurrent::Map, not a Hash
+        owner_to_pool.keys.each do |k|
+          next if k == spec_name
+
+          v = owner_to_pool[k]
+          owner_to_pool.delete(k) if v.is_a?(ConnectionPoolProxy) && v.default_pool.spec.name == spec_name
+        end
+
+        # unwrap the pool from inside a ConnectionPoolProxy
         pool = owner_to_pool[spec_name]
         owner_to_pool[spec_name] = pool.default_pool if pool.is_a?(ConnectionPoolProxy)
+
+        # now let Rails do its thing with the data type it expects
         super
       end
 
@@ -122,16 +140,21 @@ module Switchman
             else
               ancestor_pool.spec
             end
-            spec = spec.to_hash if ::Rails.version >= '5.1'
-            pool = establish_connection spec
-            pool.instance_variable_set(:@schema_cache, ancestor_pool.schema_cache) if ancestor_pool.schema_cache
-            pool
-          elsif spec_name != "primary"
+            # avoid copying "duplicate" pools that implement shard categories.
+            # they'll have a spec.name of primary, but a spec_name of something else, like unsharded
+            if spec.name == spec_name
+              pool = establish_connection(spec.to_hash)
+              pool.instance_variable_set(:@schema_cache, ancestor_pool.schema_cache) if ancestor_pool.schema_cache
+              next pool
+            end
+          end
+
+          if spec_name != "primary"
             primary_pool = retrieve_connection_pool("primary")
             if primary_pool.is_a?(ConnectionPoolProxy)
               pool = ConnectionPoolProxy.new(spec_name.to_sym, primary_pool.default_pool, @shard_connection_pools)
-              pool.schema_cache.copy_values(primary_pool.schema_cache)
-              pool
+              pool.schema_cache.copy_references(primary_pool.schema_cache)
+              owner_to_pool[spec_name] = pool
             else
               primary_pool
             end

data/lib/switchman/active_record/connection_pool.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require 'switchman/errors'
 
 module Switchman
@@ -32,10 +34,10 @@ module Switchman
         conn
       end
 
-      def connection
-        conn = super
+      def connection(switch_shard: true)
+        conn = super()
         raise NonExistentShardError if shard.new_record?
-        switch_database(conn) if conn.shard != self.shard
+        switch_database(conn) if conn.shard != self.shard && switch_shard
         conn
       end
 
@@ -47,6 +49,20 @@ module Switchman
         end
       end
 
+      def remove_shard!(shard)
+        synchronize do
+          # The shard might be currently active, so we need to update our own shard
+          if self.shard == shard
+            self.shard = Shard.default
+          end
+          # Update out any connections that may be using this shard
+          @connections.each do |conn|
+            # This will also update the connection's shard to the default shard
+            switch_database(conn) if conn.shard == shard
+          end
+        end
+      end
+
       def clear_idle_connections!(since_when)
         synchronize do
           @connections.reject! do |conn|
@@ -70,18 +86,6 @@ module Switchman
         end
 
         spec.config[:shard_name] = self.shard.name
-        case conn.adapter_name
-          when 'MySQL', 'Mysql2'
-            conn.execute("USE #{spec.config[:database]}")
-          when 'PostgreSQL'
-            if conn.schema_search_path != spec.config[:schema_search_path]
-              conn.schema_search_path = spec.config[:schema_search_path]
-            end
-          when 'SQLite'
-            # This is an artifact of the adapter modifying the path to be an absolute path when it is instantiated; just let it slide
-          else
-            raise("Cannot switch databases on same DatabaseServer with adapter type: #{conn.adapter_name}. Limit one Shard per DatabaseServer.")
-        end
         conn.shard = shard
       end