switchman 3.4.2 → 3.5.0

data/lib/switchman/active_record/database_configurations.rb

@@ -10,9 +10,9 @@ module Switchman
       def configs_for(include_replicas: false, name: nil, **)
         res = super
         if name && !include_replicas
-          return nil unless name.end_with?('primary')
+          return nil unless name.end_with?("primary")
         elsif !include_replicas
-          return res.select { |config| config.name.end_with?('primary') }
+          return res.select { |config| config.name.end_with?("primary") }
         end
         res
       end
@@ -29,19 +29,24 @@ module Switchman
         db_configs = configs.flat_map do |env_name, config|
           # It would be nice to do the auto-fallback that we want here, but we haven't
           # actually done that for years (or maybe ever) and it will be a big lift to get working
-          roles = config.keys.select { |k| config[k].is_a?(Hash) || (config[k].is_a?(Array) && config[k].all? { |ck| ck.is_a?(Hash) }) }
+          roles = config.keys.select do |k|
+            config[k].is_a?(Hash) || (config[k].is_a?(Array) && config[k].all?(Hash))
+          end
           base_config = config.except(*roles)
 
           name = "#{env_name}/primary"
-          name = 'primary' if env_name == default_env
+          name = "primary" if env_name == default_env
           base_db = build_db_config_from_raw_config(env_name, name, base_config)
           [base_db] + roles.map do |role|
-            build_db_config_from_raw_config(env_name, "#{env_name}/#{role}",
-                                            base_config.merge(config[role].is_a?(Array) ? config[role].first : config[role]))
+            build_db_config_from_raw_config(
+              env_name,
+              "#{env_name}/#{role}",
+              base_config.merge(config[role].is_a?(Array) ? config[role].first : config[role])
+            )
           end
         end
 
-        db_configs << environment_url_config(default_env, 'primary', {}) unless db_configs.find(&:for_current_env?)
+        db_configs << environment_url_config(default_env, "primary", {}) unless db_configs.find(&:for_current_env?)
 
         merge_db_environment_variables(default_env, db_configs.compact)
       end

data/lib/switchman/active_record/finder_methods.rb

@@ -49,7 +49,7 @@ module Switchman
         relation = apply_join_dependency(eager_loading: false)
         return false if ::ActiveRecord::NullRelation === relation
 
-        relation = relation.except(:select, :order).select('1 AS one').limit(1)
+        relation = relation.except(:select, :order).select("1 AS one").limit(1)
 
         case conditions
         when Array, Hash

data/lib/switchman/active_record/log_subscriber.rb

@@ -14,14 +14,14 @@ module Switchman
 
         name  = "#{payload[:name]} (#{event.duration.round(1)}ms)"
         name  = "CACHE #{name}" if payload[:cached]
-        sql   = payload[:sql].squeeze(' ')
+        sql   = payload[:sql].squeeze(" ")
         binds = nil
         shard = payload[:shard]
         shard = "  [#{shard[:database_server_id]}:#{shard[:id]} #{shard[:env]}]" if shard
 
         unless (payload[:binds] || []).empty?
           casted_params = type_casted_binds(payload[:type_casted_binds])
-          binds = '  ' + payload[:binds].zip(casted_params).map do |attr, value|
+          binds = "  " + payload[:binds].zip(casted_params).map do |attr, value|
             render_bind(attr, value)
           end.inspect
         end

data/lib/switchman/active_record/migration.rb

@@ -14,7 +14,9 @@ module Switchman
 
       def connection
         conn = super
-        ::ActiveRecord::Base.connection_pool.switch_database(conn) if conn.shard != ::ActiveRecord::Base.current_switchman_shard
+        if conn.shard != ::ActiveRecord::Base.current_switchman_shard
+          ::ActiveRecord::Base.connection_pool.switch_database(conn)
+        end
         conn
       end
     end
@@ -65,7 +67,7 @@ module Switchman
         return @migrations if instance_variable_defined?(:@migrations)
 
         migrations_cache = Thread.current[:migrations_cache] ||= {}
-        key = Digest::MD5.hexdigest(migration_files.sort.join(','))
+        key = Digest::MD5.hexdigest(migration_files.sort.join(","))
         @migrations = migrations_cache[key] ||= super
       end
 

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -5,9 +5,9 @@ module Switchman
     module PostgreSQLAdapter
       # copy/paste; use quote_local_table_name
       def create_database(name, options = {})
-        options = { encoding: 'utf8' }.merge!(options.symbolize_keys)
+        options = { encoding: "utf8" }.merge!(options.symbolize_keys)
 
-        option_string = options.sum('') do |key, value|
+        option_string = options.sum("") do |key, value|
           case key
           when :owner
             " OWNER = \"#{value}\""
@@ -24,7 +24,7 @@ module Switchman
           when :connection_limit
             " CONNECTION LIMIT = #{value}"
           else
-            ''
+            ""
           end
         end
 
@@ -37,7 +37,7 @@ module Switchman
       end
 
       def current_schemas
-        select_values('SELECT * FROM unnest(current_schemas(false))')
+        select_values("SELECT * FROM unnest(current_schemas(false))")
       end
 
       def extract_schema_qualified_name(string)
@@ -49,13 +49,13 @@ module Switchman
       # significant change: use the shard name if no explicit schema
       def quoted_scope(name = nil, type: nil)
         schema, name = extract_schema_qualified_name(name)
-        type = \
+        type =
           case type # rubocop:disable Style/HashLikeCase
-          when 'BASE TABLE'
+          when "BASE TABLE"
             "'r','p'"
-          when 'VIEW'
+          when "VIEW"
             "'v','m'"
-          when 'FOREIGN TABLE'
+          when "FOREIGN TABLE"
             "'f'"
           end
         scope = {}
@@ -67,7 +67,8 @@ module Switchman
 
       def foreign_keys(table_name)
         super.each do |fk|
-          to_table_qualified_name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(fk.to_table)
+          to_table_qualified_name =
+            ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(fk.to_table)
           fk.to_table = to_table_qualified_name.identifier if to_table_qualified_name.schema == shard.name
         end
       end
@@ -101,7 +102,7 @@ module Switchman
 
       def add_index_options(_table_name, _column_name, **)
         index, algorithm, if_not_exists = super
-        algorithm = nil if DatabaseServer.creating_new_shard && algorithm == 'CONCURRENTLY'
+        algorithm = nil if DatabaseServer.creating_new_shard && algorithm == "CONCURRENTLY"
         [index, algorithm, if_not_exists]
       end
 

data/lib/switchman/active_record/query_cache.rb

@@ -20,7 +20,7 @@ module Switchman
                 type_casted_binds: -> { type_casted_binds(binds) }
               }
               ::ActiveSupport::Notifications.instrument(
-                'sql.active_record',
+                "sql.active_record",
                 args
               )
               query_cache[sql][binds]

data/lib/switchman/active_record/query_methods.rb

@@ -84,6 +84,14 @@ module Switchman
         super(other.shard(primary_shard))
       end
 
+      # use a temp variable so that the new where clause is built before self.where_clause is read,
+      # since build_where_clause might mutate self.where_clause
+      def where!(opts, *rest)
+        new_clause = build_where_clause(opts, rest)
+        self.where_clause += new_clause
+        self
+      end
+
       protected
 
       def remove_nonlocal_primary_keys!
@@ -93,7 +101,7 @@ module Switchman
             predicate.left.relation.klass == klass &&
             (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::HomogeneousIn))
 
-          value.is_a?(Integer) && value > Shard::IDS_PER_SHARD ? [] : value
+          (value.is_a?(Integer) && value > Shard::IDS_PER_SHARD) ? [] : value
         end
         self
       end
@@ -172,7 +180,7 @@ module Switchman
 
       def sharded_primary_key?(relation, column)
         column = column.to_s
-        return column == 'id' if relation.klass == ::ActiveRecord::Base
+        return column == "id" if relation.klass == ::ActiveRecord::Base
 
         relation.klass.primary_key == column && relation.klass.integral_id?
       end
@@ -199,12 +207,11 @@ module Switchman
 
         case opts
         when String, Array
-          values = Hash === rest.first ? rest.first.values : rest
+          values = (Hash === rest.first) ? rest.first.values : rest
 
-          values.grep(ActiveRecord::Relation) do |rel|
-            # serialize subqueries against the same shard as the outer query is currently
-            # targeted to run against
-            rel.shard!(primary_shard) if rel.shard_source_value == :implicit && rel.primary_shard != primary_shard
+          if values.grep(ActiveRecord::Relation).first
+            raise "Sub-queries are not allowed as simple substitutions; " \
+                  "please build your relation with more structured methods so that Switchman is able to introspect it."
           end
 
           super
@@ -254,9 +261,10 @@ module Switchman
         send(clause_setter, new_clause)
       end
 
-      def each_transposable_predicate(predicates = nil, &block)
+      def each_transposable_predicate(predicates, &block)
         each_predicate(predicates) do |predicate|
-          if predicate.is_a?(::Arel::Nodes::Grouping)
+          case predicate
+          when ::Arel::Nodes::Grouping
             next predicate unless predicate.expr.is_a?(::Arel::Nodes::Or)
 
             or_expr = predicate.expr
@@ -267,6 +275,40 @@ module Switchman
             next predicate if new_left == old_left && new_right == old_right
 
             next predicate.class.new predicate.expr.class.new(new_left, new_right)
+          when ::Arel::Nodes::SelectStatement
+            new_cores = predicate.cores.map do |core|
+              next core unless core.is_a?(::Arel::Nodes::SelectCore) # just in case something weird is going on
+
+              new_wheres = each_transposable_predicate(core.wheres, &block)
+              new_havings = each_transposable_predicate(core.havings, &block)
+
+              next core if core.wheres == new_wheres && core.havings == new_havings
+
+              new_core = core.clone
+              new_core.wheres = new_wheres
+              new_core.havings = new_havings
+              new_core
+            end
+
+            next predicate if predicate.cores == new_cores
+
+            new_node = predicate.clone
+            new_node.instance_variable_set(:@cores, new_cores)
+            next new_node
+          when ::Arel::Nodes::Not
+            old_value = predicate.expr
+            new_value = each_transposable_predicate([old_value], &block).first
+
+            next predicate if old_value == new_value
+
+            next predicate.class.new(new_value)
+          when ::Arel::Nodes::Exists
+            old_value = predicate.expressions
+            new_value = each_transposable_predicate([old_value], &block).first
+
+            next predicate if old_value == new_value
+
+            next predicate.class.new(new_value)
           end
 
           next predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
@@ -279,54 +321,56 @@ module Switchman
         end
       end
 
-      def each_transposable_predicate_value(predicates = nil)
+      def each_transposable_predicate_value(predicates = nil, &block)
         each_transposable_predicate(predicates) do |predicate, relation, column, type|
-          each_transposable_predicate_value_cb(predicate) do |value|
+          each_transposable_predicate_value_cb(predicate, block) do |value|
             yield(value, predicate, relation, column, type)
           end
         end
       end
 
-      def each_transposable_predicate_value_cb(node, &block)
+      def each_transposable_predicate_value_cb(node, original_block, &block)
         case node
         when Array
-          node.map { |val| each_transposable_predicate_value_cb(val, &block).presence }.compact
+          node.filter_map { |val| each_transposable_predicate_value_cb(val, original_block, &block).presence }
         when ::ActiveModel::Attribute
           old_value = node.value_before_type_cast
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(node.name, new_value, node.type)
+          (old_value == new_value) ? node : node.class.new(node.name, new_value, node.type)
         when ::Arel::Nodes::And
           old_value = node.children
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(new_value)
+          (old_value == new_value) ? node : node.class.new(new_value)
         when ::Arel::Nodes::BindParam
           old_value = node.value
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(new_value)
+          (old_value == new_value) ? node : node.class.new(new_value)
         when ::Arel::Nodes::Casted
           old_value = node.value
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(new_value, node.attribute)
+          (old_value == new_value) ? node : node.class.new(new_value, node.attribute)
         when ::Arel::Nodes::HomogeneousIn
           old_value = node.values
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
           # switch to a regular In, so that Relation::WhereClause#contradiction? knows about it
           if new_value.empty?
-            klass = node.type == :in ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
+            klass = (node.type == :in) ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
             klass.new(node.attribute, new_value)
           else
-            old_value == new_value ? node : node.class.new(new_value, node.attribute, node.type)
+            (old_value == new_value) ? node : node.class.new(new_value, node.attribute, node.type)
           end
         when ::Arel::Nodes::Binary
           old_value = node.right
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(node.left, new_value)
+          (old_value == new_value) ? node : node.class.new(node.left, new_value)
+        when ::Arel::Nodes::SelectStatement
+          each_transposable_predicate_value([node], &original_block).first
         else
           yield(node)
         end
@@ -343,6 +387,8 @@ module Switchman
               Shard.current(klass.connection_class_for_self)
             elsif type == :foreign
               source_shard_for_foreign_key(relation, column)
+            else
+              primary_shard
             end
 
           transpose_predicate_value(value, current_source_shard, target_shard, type)

data/lib/switchman/active_record/relation.rb

@@ -49,7 +49,7 @@ module Switchman
       end
 
       def load(&block)
-        if !loaded? || (::Rails.version >= '7.0' && scheduled?)
+        if !loaded? || (::Rails.version >= "7.0" && scheduled?)
           @records = activate { |relation| relation.send(:exec_queries, &block) }
           @loaded = true
         end
@@ -58,7 +58,7 @@ module Switchman
       end
 
       %I[update_all delete_all].each do |method|
-        arg_params = RUBY_VERSION <= '2.8' ? '*args' : '*args, **kwargs'
+        arg_params = (RUBY_VERSION <= "2.8") ? "*args" : "*args, **kwargs"
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
           def #{method}(#{arg_params})
             result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, #{arg_params}) }
@@ -73,12 +73,16 @@ module Switchman
         loose_mode = options[:loose] && is_integer
         # loose_mode: if we don't care about getting exactly batch_size ids in between
         # don't get the max - just get the min and add batch_size so we get that many _at most_
-        values = loose_mode ? 'MIN(id)' : 'MIN(id), MAX(id)'
+        values = loose_mode ? "MIN(id)" : "MIN(id), MAX(id)"
 
         batch_size = options[:batch_size].try(:to_i) || 1000
-        quoted_primary_key = "#{klass.connection.quote_local_table_name(table_name)}.#{klass.connection.quote_column_name(primary_key)}"
-        as_id = ' AS id' unless primary_key == 'id'
-        subquery_scope = except(:select).select("#{quoted_primary_key}#{as_id}").reorder(primary_key.to_sym).limit(loose_mode ? 1 : batch_size)
+        quoted_primary_key =
+          "#{klass.connection.quote_local_table_name(table_name)}.#{klass.connection.quote_column_name(primary_key)}"
+        as_id = " AS id" unless primary_key == "id"
+        subquery_scope = except(:select)
+                         .select("#{quoted_primary_key}#{as_id}")
+                         .reorder(primary_key.to_sym)
+                         .limit(loose_mode ? 1 : batch_size)
         subquery_scope = subquery_scope.where("#{quoted_primary_key} <= ?", options[:end_at]) if options[:end_at]
 
         first_subquery_scope = if options[:start_at]
@@ -121,7 +125,9 @@ module Switchman
             relation = shard(Shard.current(klass.connection_class_for_self))
             relation.remove_nonlocal_primary_keys!
             # do a minimal query if possible
-            relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
+            if limit_value && !result_count.zero? && order_values.empty?
+              relation = relation.limit(limit_value - result_count)
+            end
 
             shard_results = relation.activate(&block)
 

data/lib/switchman/active_record/spawn_methods.rb

@@ -17,7 +17,7 @@ module Switchman
             final_shard_source_value = %i[explicit association].detect do |source_value|
               shard_source_value == source_value || rhs.shard_source_value == source_value
             end
-            raise 'unknown shard_source_value' unless final_shard_source_value
+            raise "unknown shard_source_value" unless final_shard_source_value
 
             # have to merge shard_value
             lhs_shard_value = all_shards
@@ -36,7 +36,7 @@ module Switchman
           final_shard_source_value = %i[explicit association implicit].detect do |source_value|
             shard_source_value == source_value || rhs.shard_source_value == source_value
           end
-          raise 'unknown shard_source_value' unless final_shard_source_value
+          raise "unknown shard_source_value" unless final_shard_source_value
         end
 
         [final_shard_value, final_primary_shard, final_shard_source_value]

data/lib/switchman/active_record/statement_cache.rb

@@ -4,8 +4,8 @@ module Switchman
   module ActiveRecord
     module StatementCache
       module ClassMethods
-        def create(connection, &block)
-          relation = block.call ::ActiveRecord::StatementCache::Params.new
+        def create(connection)
+          relation = yield ::ActiveRecord::StatementCache::Params.new
 
           _query_builder, binds = connection.cacheable_query(self, relation.arel)
           bind_map = ::ActiveRecord::StatementCache::BindMap.new(binds)

data/lib/switchman/active_record/tasks/database_tasks.rb

@@ -7,7 +7,7 @@ module Switchman
         def drop(*)
           super
           # no really, it's gone
-          Switchman.cache.delete('default_shard')
+          Switchman.cache.delete("default_shard")
           Shard.default(reload: true)
         end
       end

data/lib/switchman/active_record/test_fixtures.rb

@@ -12,31 +12,34 @@ module Switchman
         # Replace the one that activerecord natively uses with a switchman-optimized one
         ::ActiveSupport::Notifications.unsubscribe(@connection_subscriber)
         # Code adapted from the code in rails proper
-        @connection_subscriber = ::ActiveSupport::Notifications.subscribe('!connection.active_record') do |_, _, _, _, payload|
-          spec_name = payload[:spec_name] if payload.key?(:spec_name)
-          shard = payload[:shard] if payload.key?(:shard)
-          setup_shared_connection_pool
+        @connection_subscriber =
+          ::ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
+            spec_name = payload[:spec_name] if payload.key?(:spec_name)
+            shard = payload[:shard] if payload.key?(:shard)
+            setup_shared_connection_pool
 
-          if spec_name && !FORBIDDEN_DB_ENVS.include?(shard)
-            begin
-              connection = ::ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard: shard)
-            rescue ::ActiveRecord::ConnectionNotEstablished, ::ActiveRecord::NoDatabaseError
-              connection = nil
-            end
+            if spec_name && !FORBIDDEN_DB_ENVS.include?(shard)
+              begin
+                connection = ::ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard: shard)
+              rescue ::ActiveRecord::ConnectionNotEstablished, ::ActiveRecord::NoDatabaseError
+                connection = nil
+              end
 
-            if connection && !@fixture_connections.include?(connection)
-              connection.begin_transaction joinable: false, _lazy: false
-              connection.pool.lock_thread = true if lock_threads
-              @fixture_connections << connection
+              if connection && !@fixture_connections.include?(connection)
+                connection.begin_transaction joinable: false, _lazy: false
+                connection.pool.lock_thread = true if lock_threads
+                @fixture_connections << connection
+              end
             end
           end
-        end
       end
 
       def enlist_fixture_connections
         setup_shared_connection_pool
 
-        ::ActiveRecord::Base.connection_handler.connection_pool_list.reject { |cp| FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym) }.map(&:connection)
+        ::ActiveRecord::Base.connection_handler.connection_pool_list.reject do |cp|
+          FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym)
+        end.map(&:connection)
       end
     end
   end

data/lib/switchman/active_support/cache.rb

@@ -24,7 +24,10 @@ module Switchman
           store = super
           # can't use defined?, because it's a _ruby_ autoloaded constant,
           # so just checking that will cause it to get required
-          ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore) if store.instance_of?(ActiveSupport::Cache::RedisCacheStore) && !::ActiveSupport::Cache::RedisCacheStore.ancestors.include?(RedisCacheStore)
+          if store.instance_of?(ActiveSupport::Cache::RedisCacheStore) &&
+             !::ActiveSupport::Cache::RedisCacheStore <= RedisCacheStore
+            ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore)
+          end
           store.options[:namespace] ||= -> { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
           store
         end

data/lib/switchman/arel.rb

@@ -15,24 +15,24 @@ module Switchman
 
         def visit_Arel_Nodes_TableAlias(o, collector)
           collector = visit o.relation, collector
-          collector << ' '
+          collector << " "
           collector << quote_local_table_name(o.name)
         end
 
         def visit_Arel_Attributes_Attribute(o, collector)
           join_name = o.relation.table_alias || o.relation.name
-          collector << quote_local_table_name(join_name) << '.' << quote_column_name(o.name)
+          collector << quote_local_table_name(join_name) << "." << quote_column_name(o.name)
         end
 
         def visit_Arel_Nodes_HomogeneousIn(o, collector)
           collector.preparable = false
 
-          collector << quote_local_table_name(o.table_name) << '.' << quote_column_name(o.column_name)
+          collector << quote_local_table_name(o.table_name) << "." << quote_column_name(o.column_name)
 
           collector << if o.type == :in
-                         ' IN ('
+                         " IN ("
                        else
-                         ' NOT IN ('
+                         " NOT IN ("
                        end
 
           values = o.casted_values
@@ -43,7 +43,7 @@ module Switchman
             collector.add_binds(values, o.proc_for_binds, &bind_block)
           end
 
-          collector << ')'
+          collector << ")"
           collector
         end
 

data/lib/switchman/call_super.rb

@@ -12,7 +12,7 @@ module Switchman
       method.super_method
     end
 
-    if RUBY_VERSION <= '2.8'
+    if RUBY_VERSION <= "2.8"
       def call_super(method, above_module, *args, &block)
         super_method_above(method, above_module).call(*args, &block)
       end