switchman 1.6.1 → 2.0.9

data/lib/switchman/active_record/connection_handler.rb

@@ -1,29 +1,28 @@
+# frozen_string_literal: true
+
 require 'switchman/connection_pool_proxy'
 
 module Switchman
   module ActiveRecord
     module ConnectionHandler
-      def self.make_sharing_automagic(config)
-        key = config[:adapter] == 'postgresql' ? :schema_search_path : :database
-
+      def self.make_sharing_automagic(config, shard = Shard.current)
         # only load the shard name from the db if we have to
-        if config[key] || !config[:shard_name]
+        if !config[:shard_name]
           # we may not be able to connect to this shard yet, cause it might be an empty database server
-          shard_name = Shard.current.name rescue nil
+          shard = shard.call if shard.is_a?(Proc)
+          shard_name = shard.name rescue nil
           return unless shard_name
 
           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(owner, spec)
-        super
+      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
         # need this addition ASAP since it will be called when loading the default shard below
@@ -32,45 +31,48 @@ module Switchman
           ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter)
         end
 
-        # AR3 uses the name, AR4 uses the model
-        model = case owner
-                when String
-                  owner.constantize
-                when Class
-                  owner
-                else
-                  raise "unknown owner #{owner}"
-                end
-        pool = owner_to_pool[owner.name]
-
         first_time = !Shard.instance_variable_get(:@default)
         if first_time
           # Have to cache the default shard before we insert sharding, otherwise the first access
           # 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
-          ConnectionHandler.make_sharing_automagic(spec.config)
+          ConnectionHandler.make_sharing_automagic(config)
           ConnectionHandler.make_sharing_automagic(Shard.default.database_server.config)
 
-          ::ActiveRecord::Base.configurations[::Rails.env] = spec.instance_variable_get(:@config).stringify_keys
+          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}
 
-        proxy = ConnectionPoolProxy.new(model.shard_category,
+        @shard_connection_pools ||= { [:primary, Shard.default.database_server.shareable? ? ::Rails.env : Shard.default] => pool}
+
+        category = pool.spec.name.to_sym
+        proxy = ConnectionPoolProxy.new(category,
                                         pool,
                                         @shard_connection_pools)
-        owner_to_pool[owner.name] = proxy
-        class_to_pool.clear
+        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
 
@@ -79,7 +81,7 @@ module Switchman
         # DON'T do it if we're not the current connection handler - that means
         # we're in the middle of switching environments, and we don't want to
         # establish a connection with incorrect settings
-        if (model == ::ActiveRecord::Base || model == Shard) && self == ::ActiveRecord::Base.connection_handler && !first_time
+        if [:primary, :unsharded].include?(category) && self == ::ActiveRecord::Base.connection_handler && !first_time
           Shard.default(reload: true, with_fallback: true)
           proxy.disconnect!
         end
@@ -89,12 +91,15 @@ module Switchman
           if Shard.default.is_a?(Shard)
             DatabaseServer.all.each do |server|
               next if server == Shard.default.database_server
-              shard = server.shards.where(:name => nil).first
-              shard ||= Shard.new(:database_server => server)
-              shard.activate do
-                ConnectionHandler.make_sharing_automagic(server.config)
-                ConnectionHandler.make_sharing_automagic(proxy.current_pool.spec.config)
+
+              shard = nil
+              shard_proc = -> do
+                shard ||= server.shards.where(:name => nil).first
+                shard ||= Shard.new(:database_server => server)
+                shard
               end
+              ConnectionHandler.make_sharing_automagic(server.config, shard_proc)
+              ConnectionHandler.make_sharing_automagic(proxy.current_pool.spec.config, shard_proc)
             end
           end
           # we may have established some connections above trying to infer the shard's name.
@@ -106,60 +111,47 @@ module Switchman
         proxy
       end
 
-      def remove_connection(model)
-        uninitialize_ar(model) if owner_to_pool[model.name].is_a?(ConnectionPoolProxy)
-        result = super
-        initialize_categories
-        result
+      def remove_connection(spec_name)
+        pool = owner_to_pool[spec_name]
+        owner_to_pool[spec_name] = pool.default_pool if pool.is_a?(ConnectionPoolProxy)
+        super
       end
 
-      def pool_for(owner)
-        # copypasted from AR#ConnectionHandler other than proxy handling
-
-        owner_to_pool.fetch(owner.name) {
-          if ancestor_pool = pool_from_any_process_for(owner)
+      def retrieve_connection_pool(spec_name)
+        owner_to_pool.fetch(spec_name) do
+          if ancestor_pool = pool_from_any_process_for(spec_name)
             # A connection was established in an ancestor process that must have
             # subsequently forked. We can't reuse the connection, but we can copy
             # the specification and establish a new connection with it.
-            if ancestor_pool.is_a?(ConnectionPoolProxy)
-              establish_connection owner, ancestor_pool.default_pool.spec
+            spec = if ancestor_pool.is_a?(ConnectionPoolProxy)
+              ancestor_pool.default_pool.spec
             else
-              establish_connection owner, ancestor_pool.spec
+              ancestor_pool.spec
+            end
+            pool = establish_connection(spec.to_hash)
+            pool.instance_variable_set(:@schema_cache, ancestor_pool.schema_cache) if ancestor_pool.schema_cache
+            pool
+          elsif 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_references(primary_pool.schema_cache)
+              pool
+            else
+              primary_pool
             end
           else
-            owner_to_pool[owner.name] = nil
+            owner_to_pool[spec_name] = nil
           end
-        }
-      end
-
-      def retrieve_connection_pool(klass)
-        class_to_pool[klass.name] ||= begin
-          original_klass = klass
-          until pool = pool_for(klass)
-            klass = klass.superclass
-            break unless klass <= Base
-          end
-
-          if pool.is_a?(ConnectionPoolProxy) && pool.category != original_klass.shard_category
-            default_pool = pool.default_pool
-            pool = nil
-            class_to_pool.each_value { |p| pool = p if p.is_a?(ConnectionPoolProxy) &&
-                p.category == original_klass.shard_category &&
-                p.default_pool == default_pool }
-            pool ||= ConnectionPoolProxy.new(original_klass.shard_category, default_pool, @shard_connection_pools)
-          end
-
-          class_to_pool[original_klass.name] = pool
         end
       end
 
       def clear_idle_connections!(since_when)
-        # TODO in rails 4.2+ s/connection_pools.values/connection_pool_list/
-        connection_pools.values.each{ |pool| pool.clear_idle_connections!(since_when) }
+        connection_pool_list.each{ |pool| pool.clear_idle_connections!(since_when) }
       end
 
       def switchman_connection_pool_proxies
-        class_to_pool.values.uniq
+        owner_to_pool.values.uniq.select{|p| p.is_a?(ConnectionPoolProxy)}
       end
 
       private

data/lib/switchman/active_record/connection_pool.rb

@@ -1,14 +1,16 @@
+# frozen_string_literal: true
+
 require 'switchman/errors'
 
 module Switchman
   module ActiveRecord
     module ConnectionPool
       def shard
-        Thread.current["#{object_id}_shard".to_sym] || Shard.default
+        Thread.current[tls_key] || Shard.default
       end
 
       def shard=(value)
-        Thread.current["#{object_id}_shard".to_sym] = value
+        Thread.current[tls_key] = value
       end
 
       def default_schema
@@ -32,19 +34,14 @@ 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
 
-      def release_connection(with_id = nil)
-        with_id ||= if ::Rails.version >= '5'
-                      Thread.current
-                    else
-                      current_connection_id
-                    end
+      def release_connection(with_id = Thread.current)
         super(with_id)
 
         if spec.config[:idle_timeout]
@@ -52,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|
@@ -75,20 +86,14 @@ 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
+
+      private
+
+      def tls_key
+        "#{object_id}_shard".to_sym
+      end
     end
   end
 end

data/lib/switchman/active_record/finder_methods.rb

@@ -1,9 +1,10 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module FinderMethods
-      def find_one(id, call_super: false)
+      def find_one(id)
         return super(id) unless klass.integral_id?
-        return super(id) if call_super
 
         if shard_source_value != :implicit
           current_shard = Shard.current(klass.shard_category)
@@ -14,7 +15,7 @@ module Switchman
             # skip the shard if the object can't be on it. unless we're only looking at one shard;
             # we might be expecting a shadow object
             next if current_id > Shard::IDS_PER_SHARD && self.all_shards.length > 1
-            relation.send(:find_one, current_id, call_super: true)
+            relation.call_super(:find_one, FinderMethods, current_id)
           end
           if result.is_a?(Array)
             result = result.first
@@ -26,24 +27,18 @@ module Switchman
 
         local_id, shard = Shard.local_id_for(id)
         if shard
-          if ::Rails.version < '4.2'
-            # find_one uses binds, so we can't depend on QueryMethods
-            # catching it
-            begin
-              old_shard_value = shard_value
-              self.shard_value = shard
-              super(local_id)
-            ensure
-              self.shard_value = old_shard_value
-            end
-          else
-            shard.activate { super(local_id) }
-          end
+          shard.activate { super(local_id) }
         else
           super(id)
         end
       end
 
+      def find_some_ordered(ids)
+        current_shard = Shard.current(klass.shard_category)
+        ids = ids.map{|id| Shard.relative_id_for(id, current_shard, current_shard)}
+        super(ids)
+      end
+
       def find_or_instantiator_by_attributes(match, attributes, *args)
         primary_shard.activate { super }
       end
@@ -52,13 +47,10 @@ module Switchman
         conditions = conditions.id if ::ActiveRecord::Base === conditions
         return false if !conditions
 
-        if ::Rails.version >= '4.1'
-          relation = apply_join_dependency(self, construct_join_dependency)
-          return false if ::ActiveRecord::NullRelation === relation
-        else
-          join_dependency = construct_join_dependency_for_association_find
-          relation = construct_relation_for_association_find(join_dependency)
-        end
+        relation = ::Rails.version >= "5.2" ?
+          apply_join_dependency(eager_loading: false) :
+          apply_join_dependency(self, construct_join_dependency)
+        return false if ::ActiveRecord::NullRelation === relation
 
         relation = relation.except(:select, :order).select("1 AS one").limit(1)
 
@@ -69,12 +61,11 @@ module Switchman
           relation = relation.where(table[primary_key].eq(conditions)) if conditions != :none
         end
 
-        args = [relation, "#{name} Exists"]
-        args << relation.bind_values if ::Rails.version >= '4.1'
-        relation.activate { return true if connection.select_value(*args) }
-        false
-      rescue
-        raise if ::Rails.version >= '4.1' || !(::ActiveRecord::ThrowResult === $!)
+        relation.activate do |shard_rel|
+          return true if ::Rails.version >= "5.2" ?
+            connection.select_value(shard_rel.arel, "#{name} Exists") :
+            connection.select_value(shard_rel, "#{name} Exists", shard_rel.bound_attributes)
+        end
         false
       end
     end

data/lib/switchman/active_record/log_subscriber.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module LogSubscriber
@@ -8,35 +10,28 @@ module Switchman
 
         payload = event.payload
 
-        return if 'SCHEMA'.freeze == payload[:name]
+        return if ::ActiveRecord::LogSubscriber::IGNORE_PAYLOAD_NAMES.include?(payload[:name])
 
-        name  = '%s (%.1fms)'.freeze % [payload[:name], event.duration]
+        name  = "#{payload[:name]} (#{event.duration.round(1)}ms)"
+        name  = "CACHE #{name}" if payload[:cached]
         sql   = payload[:sql].squeeze(' '.freeze)
         binds = nil
         shard = payload[:shard]
         shard = "  [#{shard[:database_server_id]}:#{shard[:id]} #{shard[:env]}]" if shard
 
         unless (payload[:binds] || []).empty?
-          binds = "  " + payload[:binds].map { |col,v|
-            if col
-              [col.name, v]
-            else
-              [nil, v]
-            end
+          use_old_format = (::Rails.version < '5.1.5')
+          args = use_old_format ?
+            [payload[:binds], payload[:type_casted_binds]] :
+            [payload[:type_casted_binds]]
+          casted_params = type_casted_binds(*args)
+          binds = "  " + payload[:binds].zip(casted_params).map { |attr, value|
+            render_bind(attr, value)
           }.inspect
         end
 
-        if ::Rails.version >= '5'
-          name = colorize_payload_name(name, payload[:name])
-          sql  = color(sql, sql_color(sql), true)
-        else
-          if odd?
-            name = color(name, self.class::CYAN, true)
-            sql  = color(sql, nil, true)
-          else
-            name = color(name, self.class::MAGENTA, true)
-          end
-        end
+        name = colorize_payload_name(name, payload[:name])
+        sql  = color(sql, sql_color(sql), true)
 
         debug "  #{name}  #{sql}#{binds}#{shard}"
       end

data/lib/switchman/active_record/migration.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module Migration
+      module Compatibility
+        module V5_0
+          def create_table(*args, **options)
+            unless options.key?(:id)
+              options[:id] = :bigserial
+            end
+            if block_given?
+              super do |td|
+                yield td
+              end
+            else
+              super
+            end
+          end
+        end
+      end
+
+      def connection
+        conn = super
+        if conn.shard != ::ActiveRecord::Base.connection_pool.current_pool.shard
+          ::ActiveRecord::Base.connection_pool.current_pool.switch_database(conn)
+        end
+        conn
+      end
+    end
+
+    module Migrator
+      # significant change: hash shard id, not database name
+      def generate_migrator_advisory_lock_id
+        shard_name_hash = Zlib.crc32(Shard.current.name)
+        ::ActiveRecord::Migrator::MIGRATOR_SALT * shard_name_hash
+      end
+
+      if ::Rails.version >= '6.0'
+        # copy/paste from Rails 6.1
+        def with_advisory_lock
+          lock_id = generate_migrator_advisory_lock_id
+
+          with_advisory_lock_connection do |connection|
+            got_lock = connection.get_advisory_lock(lock_id)
+            raise ::ActiveRecord::ConcurrentMigrationError unless got_lock
+            load_migrated # reload schema_migrations to be sure it wasn't changed by another process before we got the lock
+            yield
+          ensure
+            if got_lock && !connection.release_advisory_lock(lock_id)
+              raise ::ActiveRecord::ConcurrentMigrationError.new(
+                ::ActiveRecord::ConcurrentMigrationError::RELEASE_LOCK_FAILED_MESSAGE
+              )
+            end
+          end
+        end
+
+        # significant change: strip out prefer_secondary from config
+        def with_advisory_lock_connection
+          pool = ::ActiveRecord::ConnectionAdapters::ConnectionHandler.new.establish_connection(
+            ::ActiveRecord::Base.connection_config.except(:prefer_secondary)
+          )
+
+          pool.with_connection { |connection| yield(connection) }
+        ensure
+          pool&.disconnect!
+        end
+      end
+    end
+
+    module MigrationContext
+      def migrations
+        return @migrations if instance_variable_defined?(:@migrations)
+        migrations_cache = Thread.current[:migrations_cache] ||= {}
+        key = Digest::MD5.hexdigest(migration_files.sort.join(','))
+        @migrations = migrations_cache[key] ||= super
+      end
+    end
+  end
+end

data/lib/switchman/active_record/model_schema.rb

@@ -1,10 +1,12 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module ModelSchema
       module ClassMethods
         def quoted_table_name
           @quoted_table_name ||= {}
-          @quoted_table_name[Shard.current.id] ||= connection.quote_table_name(table_name)
+          @quoted_table_name[Shard.current(shard_category).id] ||= connection.quote_table_name(table_name)
         end
       end
     end

data/lib/switchman/active_record/persistence.rb

@@ -1,10 +1,18 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Persistence
       # touch reads the id attribute directly, so it's not relative to the current shard
-      def touch(*)
+      def touch(*, **)
         shard.activate(self.class.shard_category) { super }
       end
+
+      if ::Rails.version >= '5.2'
+        def update_columns(*)
+          shard.activate(self.class.shard_category) { super }
+        end
+      end
     end
   end
 end