switchman 3.0.2 → 3.1.0

data/lib/switchman/active_record/relation.rb

@@ -9,13 +9,13 @@ module Switchman
 
       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
 
@@ -29,41 +29,38 @@ module Switchman
       end
 
       def new(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def create(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def create!(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        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) }
       end
 
-      def records
-        return @records if loaded?
-
-        results = activate { |relation| relation.call_super(:records, Relation) }
-        case shard_value
-        when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
-          @records = results
+      def load(&block)
+        if !loaded? || (::Rails.version >= '7.0' && 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) }
+            result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
             result = result.sum if result.is_a?(Array)
             result
           end
@@ -94,7 +91,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 +100,28 @@ module Switchman
         end
       end
 
-      def activate(&block)
+      def activate(unordered: false, &block)
         shards = all_shards
         if Array === shards && shards.length == 1
-          if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_classes)
+          if 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), :to_a)
             # do a minimal query if possible
             relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
 
             shard_results = relation.activate(&block)
 
-            if shard_results.present?
+            if 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 +145,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/statement_cache.rb

@@ -33,12 +33,12 @@ module Switchman
           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)
 
-        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)
         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: limit)
         super
       end
     end

data/lib/switchman/active_record/test_fixtures.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module TestFixtures
+      FORBIDDEN_DB_ENVS = %i[development production].freeze
+      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[: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 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
+
+      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)
+      end
+    end
+  end
+end

data/lib/switchman/active_support/cache.rb

@@ -4,6 +4,22 @@ 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,

data/lib/switchman/arel.rb

@@ -11,22 +11,44 @@ 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_TableAlias(o, collector)
           collector = visit o.relation, 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
+
+        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 << if o.type == :in
+                         ' IN ('
+                       else
+                         ' NOT IN ('
+                       end
+
+          values = o.casted_values
+
+          if values.empty?
+            collector << @connection.quote(nil)
+          else
+            collector.add_binds(values, o.proc_for_binds, &bind_block)
+          end
+
+          collector << ')'
+          collector
         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/database_server.rb

@@ -8,15 +8,12 @@ module Switchman
 
     class << self
       attr_accessor :creating_new_shard
+      attr_reader :all_roles
 
       def all
         database_servers.values
       end
 
-      def all_roles
-        @all_roles ||= all.map(&:roles).flatten.uniq
-      end
-
       def find(id_or_all)
         return all if id_or_all == :all
         return id_or_all.map { |id| database_servers[id || ::Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
@@ -38,7 +35,7 @@ module Switchman
         database_servers[server.id] = server
         ::ActiveRecord::Base.configurations.configurations <<
           ::ActiveRecord::DatabaseConfigurations::HashConfig.new(::Rails.env, "#{server.id}/primary", settings)
-        Shard.send(:initialize_sharding)
+        Shard.send(:configure_connects_to)
         server
       end
 
@@ -49,31 +46,42 @@ module Switchman
         servers[rand(servers.length)]
       end
 
+      def guard_servers
+        all.each { |db| db.guard! if db.config[:prefer_secondary] }
+      end
+
       private
 
       def reference_role(role)
         return if all_roles.include?(role)
 
         @all_roles << role
-        Shard.send(:initialize_sharding)
+        Shard.send(:configure_connects_to)
       end
 
       def database_servers
-        unless @database_servers
+        if !@database_servers || @database_servers.empty?
           @database_servers = {}.with_indifferent_access
+          roles = []
           ::ActiveRecord::Base.configurations.configurations.each do |config|
             if config.name.include?('/')
               name, role = config.name.split('/')
             else
               name, role = config.env_name, config.name
             end
+            role = role.to_sym
 
-            if role == 'primary'
+            roles << role
+            if role == :primary
               @database_servers[name] = DatabaseServer.new(config.env_name, config.configuration_hash)
             else
               @database_servers[name].roles << role
             end
           end
+          # Do this after so that all database servers for all roles are established and we won't prematurely
+          # configure a connection for the wrong role
+          @all_roles = roles.uniq
+          Shard.send(:configure_connects_to)
         end
         @database_servers
       end
@@ -89,13 +97,13 @@ module Switchman
     end
 
     def connects_to_hash
-      self.class.all_roles.map do |role|
+      self.class.all_roles.to_h do |role|
         config_role = role
         config_role = :primary unless roles.include?(role)
         config_name = :"#{id}/#{config_role}"
         config_name = :primary if id == ::Rails.env && config_role == :primary
         [role.to_sym, config_name]
-      end.to_h
+      end
     end
 
     def destroy
@@ -129,27 +137,27 @@ module Switchman
         end
     end
 
-    def guard_rail_environment
-      @guard_rail_environment || ::GuardRail.environment
-    end
-
     # locks this db to a specific environment, except for
     # when doing writes (then it falls back to the current
     # value of GuardRail.environment)
     def guard!(environment = :secondary)
-      @guard_rail_environment = environment
+      DatabaseServer.send(:reference_role, environment)
+      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => environment }, klasses: [::ActiveRecord::Base] }
     end
 
     def unguard!
-      @guard_rail_environment = nil
+      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => :_switchman_inherit }, klasses: [::ActiveRecord::Base] }
     end
 
     def unguard
-      old_env = @guard_rail_environment
-      unguard!
-      yield
-    ensure
-      guard!(old_env)
+      return yield unless ::ActiveRecord::Base.role_overriden?(id.to_sym)
+
+      begin
+        unguard!
+        yield
+      ensure
+        ::ActiveRecord::Base.connected_to_stack.pop
+      end
     end
 
     def shards
@@ -186,58 +194,56 @@ module Switchman
 
       name ||= "#{config[:database]}_shard_#{id}"
 
+      schema_already_existed = false
+      shard = nil
       Shard.connection.transaction do
-        shard = Shard.create!(id: id,
-                              name: name,
-                              database_server_id: self.id)
-        schema_already_existed = false
-
-        begin
-          self.class.creating_new_shard = true
-          DatabaseServer.send(:reference_role, :deploy)
-          ::ActiveRecord::Base.connected_to(shard: self.id.to_sym, role: :deploy) do
-            if create_statement
-              if ::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}")
-                schema_already_existed = true
-                raise 'This schema already exists; cannot overwrite'
-              end
-              Array(create_statement.call).each do |stmt|
-                ::ActiveRecord::Base.connection.execute(stmt)
-              end
+        self.class.creating_new_shard = true
+        DatabaseServer.send(:reference_role, :deploy)
+        ::ActiveRecord::Base.connected_to(shard: self.id.to_sym, role: :deploy) do
+          shard = Shard.create!(id: id,
+                                name: name,
+                                database_server_id: self.id)
+          if create_statement
+            if ::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}")
+              schema_already_existed = true
+              raise 'This schema already exists; cannot overwrite'
             end
-            if config[:adapter] == 'postgresql'
-              old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor do
-              end
+            Array(create_statement.call).each do |stmt|
+              ::ActiveRecord::Base.connection.execute(stmt)
+            end
+          end
+          if config[:adapter] == 'postgresql'
+            old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor do
             end
-            old_verbose = ::ActiveRecord::Migration.verbose
-            ::ActiveRecord::Migration.verbose = false
-
-            unless schema == false
-              shard.activate(*Shard.sharded_models) do
-                reset_column_information
-
-                ::ActiveRecord::Base.connection.transaction(requires_new: true) do
-                  ::ActiveRecord::Base.connection.migration_context.migrate
-                end
-                reset_column_information
-                ::ActiveRecord::Base.descendants.reject do |m|
-                  m <= UnshardedRecord || !m.table_exists?
-                end.each(&:define_attribute_methods)
+          end
+          old_verbose = ::ActiveRecord::Migration.verbose
+          ::ActiveRecord::Migration.verbose = false
+
+          unless schema == false
+            shard.activate do
+              reset_column_information
+
+              ::ActiveRecord::Base.connection.transaction(requires_new: true) do
+                ::ActiveRecord::Base.connection.migration_context.migrate
               end
+              reset_column_information
+              ::ActiveRecord::Base.descendants.reject do |m|
+                m <= UnshardedRecord || !m.table_exists?
+              end.each(&:define_attribute_methods)
             end
-          ensure
-            ::ActiveRecord::Migration.verbose = old_verbose
-            ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
           end
-          shard
-        rescue
-          shard.destroy
-          shard.drop_database rescue nil unless schema_already_existed
-          reset_column_information unless schema == false rescue nil
-          raise
         ensure
-          self.class.creating_new_shard = false
+          ::ActiveRecord::Migration.verbose = old_verbose
+          ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
         end
+        shard
+      rescue
+        shard&.destroy
+        shard&.drop_database rescue nil unless schema_already_existed
+        reset_column_information unless schema == false rescue nil
+        raise
+      ensure
+        self.class.creating_new_shard = false
       end
     end
 

data/lib/switchman/default_shard.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require 'switchman/database_server'
-
 module Switchman
   class DefaultShard
     def id