switchman 1.5.21 → 2.1.5

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"
@@ -53,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")]
@@ -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)

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,65 @@ 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
-      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
 
-      def pool_for(owner)
-        # copypasted from AR#ConnectionHandler other than proxy handling
+          v = owner_to_pool[k]
+          owner_to_pool.delete(k) if v.is_a?(ConnectionPoolProxy) && v.default_pool.spec.name == spec_name
+        end
 
-        owner_to_pool.fetch(owner.name) {
-          if ancestor_pool = pool_from_any_process_for(owner)
+        # 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
+
+      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
+            # 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
-          else
-            owner_to_pool[owner.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)
+          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_references(primary_pool.schema_cache)
+              owner_to_pool[spec_name] = pool
+            else
+              primary_pool
+            end
+          else
+            owner_to_pool[spec_name] = nil
           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,31 +1,44 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module FinderMethods
       def find_one(id)
+        return super(id) unless klass.integral_id?
+
         if shard_source_value != :implicit
-          return self.activate { super(Shard.relative_id_for(id, Shard.current(klass.shard_category), primary_shard)) }
+          current_shard = Shard.current(klass.shard_category)
+          result = self.activate do |relation, shard|
+            current_id = Shard.relative_id_for(id, current_shard, shard)
+            # current_id will be nil for non-integral id
+            next unless current_id
+            # 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.call_super(:find_one, FinderMethods, current_id)
+          end
+          if result.is_a?(Array)
+            result = result.first
+          end
+          # we may have skipped all shards
+          raise_record_not_found_exception!(id, 0, 1) unless result
+          return result
         end
 
         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
+          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
@@ -34,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)
 
@@ -51,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'
-        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