switchman 3.0.2 → 4.2.5

data/lib/switchman/active_record/relation.rb

@@ -4,18 +4,18 @@ module Switchman
   module ActiveRecord
     module Relation
       def self.prepended(klass)
-        klass::SINGLE_VALUE_METHODS.concat %i[shard shard_source]
+        klass::SINGLE_VALUE_METHODS.push(:shard, :shard_source)
       end
 
       def initialize(*, **)
         super
-        self.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
+        self.shard_value = Shard.current(klass ? klass.connection_class_for_self : :primary) unless shard_value
         self.shard_source_value = :implicit unless shard_source_value
       end
 
       def clone
         result = super
-        result.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
+        result.shard_value = Shard.current(klass ? klass.connection_class_for_self : :primary) unless shard_value
         result
       end
 
@@ -28,59 +28,119 @@ module Switchman
         relation
       end
 
-      def new(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+      def new(*, &)
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
-      def create(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+      def create(*, &)
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
-      def create!(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+      def create!(*, &)
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def to_sql
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
-      def explain
-        activate { |relation| relation.call_super(:explain, Relation) }
+      if ::Rails.version > "7.1.2"
+        def transaction(...)
+          primary_shard.activate(klass.connection_class_for_self) { super }
+        end
       end
 
-      def records
-        return @records if loaded?
+      def explain(*args)
+        activate { |relation| relation.call_super(:explain, Relation, *args) }
+      end
 
-        results = activate { |relation| relation.call_super(:records, Relation) }
-        case shard_value
-        when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
-          @records = results
+      def load(&block)
+        if !loaded? || scheduled?
+          @records = activate { |relation| relation.send(:exec_queries, &block) }
           @loaded = true
         end
-        results
+
+        self
       end
 
       %I[update_all delete_all].each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
-          def #{method}(*args)
-            result = self.activate { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
+          def #{method}(*args, **kwargs)
+            result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, *args, **kwargs) }
             result = result.sum if result.is_a?(Array)
             result
           end
         RUBY
       end
 
+      # https://github.com/rails/rails/commit/ed2c15b52450ff927a05629f031376f25b670335
+      # once the minimum version is Rails 7.2, we can drop this separate module
+      module InsertUpsertAll
+        %w[insert_all upsert_all].each do |method|
+          class_eval <<-RUBY, __FILE__, __LINE__ + 1
+            def #{method}(attributes, returning: nil, **)
+              scope = self != ::ActiveRecord::Base && current_scope
+              if (target_shard = scope&.primary_shard) == (current_shard = Shard.current(connection_class_for_self))
+                scope = nil
+              end
+              if scope
+                dupped = false
+                attributes.each_with_index do |hash, i|
+                  if dupped || hash.any? { |k, v| sharded_column?(k) }
+                    unless dupped
+                      attributes = attributes.dup
+                      dupped = true
+                    end
+                    attributes[i] = hash.to_h do |k, v|
+                      if sharded_column?(k)
+                        [k, Shard.relative_id_for(v, current_shard, target_shard)]
+                      else
+                        [k, v]
+                      end
+                    end
+                  end
+                end
+              end
+
+              if scope
+                scope.activate do
+                  db = Shard.current(connection_class_for_self).database_server
+                  result = db.unguard { super }
+                  if result&.columns&.any? { |c| sharded_column?(c) }
+                    transposed_rows = result.rows.map do |row|
+                      row.map.with_index do |value, i|
+                        sharded_column?(result.columns[i]) ? Shard.relative_id_for(value, target_shard, current_shard) : value
+                      end
+                    end
+                    result = ::ActiveRecord::Result.new(result.columns, transposed_rows, result.column_types)
+                  end
+
+                  result
+                end
+              else
+                db = Shard.current(connection_class_for_self).database_server
+                db.unguard { super }
+              end
+            end
+          RUBY
+        end
+      end
+
       def find_ids_in_ranges(options = {})
         is_integer = columns_hash[primary_key.to_s].type == :integer
         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]
@@ -94,7 +154,7 @@ module Switchman
 
         while ids.first.present?
           ids.map!(&:to_i) if is_integer
-          ids << ids.first + batch_size if loose_mode
+          ids << (ids.first + batch_size) if loose_mode
 
           yield(*ids)
           last_value = ids.last
@@ -103,28 +163,39 @@ module Switchman
         end
       end
 
-      def activate(&block)
+      def activate(count: false, unordered: false, &block)
         shards = all_shards
         if Array === shards && shards.length == 1
-          if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_classes)
+          if !loaded? && shard_value != shards.first
+            shard(shards.first).activate(&block)
+          elsif shards.first == DefaultShard || shards.first == Shard.current(klass.connection_class_for_self)
             yield(self, shards.first)
           else
-            shards.first.activate(klass.connection_classes) { yield(self, shards.first) }
+            shards.first.activate(klass.connection_class_for_self) { yield(self, shards.first) }
           end
         else
           result_count = 0
           can_order = false
-          result = Shard.with_each_shard(shards, [klass.connection_classes]) do
+          result = Shard.with_each_shard(shards, [klass.connection_class_for_self]) do
             # don't even query other shards if we're already past the limit
             next if limit_value && result_count >= limit_value && order_values.empty?
 
-            relation = shard(Shard.current(klass.connection_classes), :to_a)
+            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)
 
-            if shard_results.present?
+            if shard_results.present? && count
+              unless shard_results.is_a?(Integer)
+                raise "expected integer result for count, got #{shard_results.class.name}"
+              end
+
+              result_count += shard_results
+            elsif shard_results.present? && !unordered
               can_order ||= can_order_cross_shard_results? unless order_values.empty?
               raise OrderOnMultiShardQuery if !can_order && !order_values.empty? && result_count.positive?
 
@@ -148,8 +219,12 @@ module Switchman
         results.sort! do |l, r|
           result = 0
           order_values.each do |ov|
-            a = l.attribute(ov.expr.name)
-            b = r.attribute(ov.expr.name)
+            if l.is_a?(::ActiveRecord::Base)
+              a = l.attribute(ov.expr.name)
+              b = r.attribute(ov.expr.name)
+            else
+              a, b = l, r
+            end
             next if a == b
 
             if a.nil? || b.nil?

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]
@@ -62,16 +62,12 @@ module Switchman
         if primary_shard != final_primary_shard && rhs.primary_shard != final_primary_shard
           shard!(final_primary_shard)
           rhs = rhs.shard(final_primary_shard)
-          super(rhs)
         elsif primary_shard != final_primary_shard
           shard!(final_primary_shard)
-          super(rhs)
         elsif rhs.primary_shard != final_primary_shard
           rhs = rhs.shard(final_primary_shard)
-          super(rhs)
-        else
-          super
         end
+        super
 
         self.shard_value = final_shard_value
         self.shard_source_value = 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)
@@ -25,22 +25,22 @@ module Switchman
       # we can make some assumptions about the shard source
       # (e.g. infer from the primary key or use the current shard)
 
-      def execute(*args)
+      def execute(*args, &block)
         params, connection = args
         klass = @klass
         target_shard = nil
-        if (primary_index = bind_map.primary_value_index)
+        if (primary_index = @bind_map.primary_value_index)
           primary_value = params[primary_index]
           target_shard = Shard.local_id_for(primary_value)[1]
         end
-        current_shard = Shard.current(klass.connection_classes)
+        current_shard = Shard.current(klass.connection_class_for_self)
         target_shard ||= current_shard
 
-        bind_values = bind_map.bind(params, current_shard, target_shard)
+        bind_values = @bind_map.bind(params, current_shard, target_shard)
 
-        target_shard.activate(klass.connection_classes) do
+        target_shard.activate(klass.connection_class_for_self) do
           sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)
-          klass.find_by_sql(sql, bind_values)
+          klass.find_by_sql(sql, bind_values, &block)
         end
       end
 
@@ -66,7 +66,11 @@ module Switchman
 
         def primary_value_index
           primary_ba_index = @bound_attributes.index do |ba|
-            ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.primary
+            if ba.value.is_a?(::ActiveRecord::StatementCache::Substitute)
+              ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.primary
+            else
+              false
+            end
           end
           @indexes.index(primary_ba_index) if primary_ba_index
         end

data/lib/switchman/active_record/table_definition.rb

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module TableDefinition
       def column(name, type, limit: nil, **)
-        Engine.foreign_key_check(name, type, limit: limit)
+        Switchman.foreign_key_check(name, type, limit:)
         super
       end
     end

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

@@ -7,9 +7,14 @@ 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
+
+        def raise_for_multi_db(*)
+          # ignore; Switchman doesn't use namespaced tasks for multiple shards; it uses
+          # environment variables to filter which shards you want to target
+        end
       end
     end
   end

data/lib/switchman/active_record/test_fixtures.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module TestFixtures
+      FORBIDDEN_DB_ENVS = %i[development production].freeze
+
+      if ::Rails.version < "7.2"
+        def setup_fixtures(config = ::ActiveRecord::Base)
+          super
+          return unless run_in_transaction?
+
+          # 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[:connection_name] if payload.key?(:connection_name))
+              shard = payload[:shard] if payload.key?(:shard)
+
+              if spec_name && !FORBIDDEN_DB_ENVS.include?(shard)
+                begin
+                  connection = ::ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard:)
+                  connection.connect! # eagerly validate the connection
+                rescue ::ActiveRecord::ConnectionNotEstablished
+                  connection = nil
+                end
+
+                if connection
+                  setup_shared_connection_pool
+                  unless @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(:primary).reject do |cp|
+            FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym)
+          end.map(&:connection)
+        end
+      else
+        def setup_transactional_fixtures
+          setup_shared_connection_pool
+
+          # Begin transactions for connections already established
+          # INST: :writing -> :primary
+          @fixture_connection_pools = ::ActiveRecord::Base.connection_handler.connection_pool_list(:primary)
+          # INST: filter by FORBIDDEN_DB_ENVS
+          @fixture_connection_pools = @fixture_connection_pools.reject do |cp|
+            FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym)
+          end
+
+          @fixture_connection_pools.each do |pool|
+            pool.pin_connection!(lock_threads)
+            pool.lease_connection
+          end
+
+          # When connections are established in the future, begin a transaction too
+          @connection_subscriber = ::ActiveSupport::Notifications
+                                   .subscribe("!connection.active_record") do |_, _, _, _, payload|
+            connection_name = payload[:connection_name] if payload.key?(:connection_name)
+            shard = payload[:shard] if payload.key?(:shard)
+
+            # INST: filter by FORBIDDEN_DB_ENVS
+            if connection_name && !FORBIDDEN_DB_ENVS.include?(shard)
+              pool = ::ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_name, shard:)
+              if pool
+                setup_shared_connection_pool
+
+                unless @fixture_connection_pools.include?(pool)
+                  pool.pin_connection!(lock_threads)
+                  pool.lease_connection
+                  @fixture_connection_pools << pool
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/switchman/active_support/cache.rb

@@ -4,11 +4,32 @@ module Switchman
   module ActiveSupport
     module Cache
       module ClassMethods
+        def lookup_stores(cache_store_config)
+          result = {}
+          cache_store_config.each do |key, value|
+            next if value.is_a?(String)
+
+            result[key] = ::ActiveSupport::Cache.lookup_store(value)
+          end
+
+          cache_store_config.each do |key, value| # rubocop:disable Style/CombinableLoops
+            next unless value.is_a?(String)
+
+            result[key] = result[value]
+          end
+          result
+        end
+
         def lookup_store(*store_options)
           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)
+          # must use the string name, otherwise it will try to auto-load the constant
+          # and we don't want to require redis in this file (since it's not a hard dependency)
+          # rubocop:disable Style/ClassEqualityComparison
+          if store.class.name == "ActiveSupport::Cache::RedisCacheStore" &&
+             !(::ActiveSupport::Cache::RedisCacheStore <= RedisCacheStore)
+            ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore)
+          end
+          # rubocop:enable Style/ClassEqualityComparison
           store.options[:namespace] ||= -> { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
           store
         end
@@ -17,7 +38,7 @@ module Switchman
       module RedisCacheStore
         def clear(namespace: nil, **)
           # RedisCacheStore tries to be smart and only clear the cache under your namespace, if you have one set
-          # unfortunately, it uses the keys command, which is extraordinarily inefficient in a large redis instance
+          # unfortunately, it doesn't work using redis clustering because of the way redis keys are distributed
           # fortunately, we can assume we control the entire instance, because we set up the namespacing, so just
           # always unset it temporarily for clear calls
           namespace = nil # rubocop:disable Lint/ShadowedArgument

data/lib/switchman/arel.rb

@@ -11,22 +11,35 @@ module Switchman
     module Visitors
       module ToSql
         # rubocop:disable Naming/MethodName
+        # rubocop:disable Naming/MethodParameterName
 
-        def visit_Arel_Nodes_TableAlias(*args)
-          o, collector = args
+        def visit_Arel_Nodes_Cte(o, collector)
+          collector << quote_local_table_name(o.name)
+          collector << " AS "
+
+          case o.materialized
+          when true
+            collector << "MATERIALIZED "
+          when false
+            collector << "NOT MATERIALIZED "
+          end
+
+          visit o.relation, collector
+        end
+
+        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(*args)
-          o = args.first
+        def visit_Arel_Attributes_Attribute(o, collector)
           join_name = o.relation.table_alias || o.relation.name
-          result = "#{quote_local_table_name join_name}.#{quote_column_name o.name}"
-          args.last << result
+          collector << quote_local_table_name(join_name) << "." << quote_column_name(o.name)
         end
 
         # rubocop:enable Naming/MethodName
+        # rubocop:enable Naming/MethodParameterName
 
         def quote_local_table_name(name)
           return name if ::Arel::Nodes::SqlLiteral === name

data/lib/switchman/call_super.rb

@@ -12,8 +12,8 @@ module Switchman
       method.super_method
     end
 
-    def call_super(method, above_module, *args, &block)
-      super_method_above(method, above_module).call(*args, &block)
+    def call_super(method, above_module, ...)
+      super_method_above(method, above_module).call(...)
     end
   end
 end