switchman 3.0.14 → 3.5.20

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

@@ -12,31 +12,41 @@ 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 = if ::Rails.version < "7.1"
+                          payload[:spec_name] if payload.key?(:spec_name)
+                        elsif payload.key?(:connection_name)
+                          payload[:connection_name]
+                        end
+            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: 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)
+                connection.connect! if ::Rails.version >= "7.1" # eagerly validate the connection
+              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
+                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
       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(:primary).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

@@ -4,11 +4,30 @@ 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)
+          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

@@ -13,38 +13,54 @@ module Switchman
         # rubocop:disable Naming/MethodName
         # rubocop:disable Naming/MethodParameterName
 
+        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(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
+        if ::Rails.version < "7.1"
+          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 ('
-                       else
-                         ' NOT IN ('
-                       end
+            collector << if o.type == :in
+                           " IN ("
+                         else
+                           " NOT IN ("
+                         end
 
-          values = o.casted_values
+            values = o.casted_values
 
-          if values.empty?
-            collector << @connection.quote(nil)
-          else
-            collector.add_binds(values, o.proc_for_binds, &bind_block)
-          end
+            if values.empty?
+              collector << @connection.quote(nil)
+            else
+              collector.add_binds(values, o.proc_for_binds, &bind_block)
+            end
 
-          collector << ')'
-          collector
+            collector << ")"
+            collector
+          end
         end
 
         # rubocop:enable Naming/MethodName

data/lib/switchman/call_super.rb

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

data/lib/switchman/database_server.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'securerandom'
+require "securerandom"
 
 module Switchman
   class DatabaseServer
@@ -8,24 +8,25 @@ module Switchman
 
     class << self
       attr_accessor :creating_new_shard
+      attr_reader :all_roles
+
+      include Enumerable
+
+      delegate :each, to: :all
 
       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)
+        return id_or_all.filter_map { |id| database_servers[id || ::Rails.env] }.uniq if id_or_all.is_a?(Array)
 
         database_servers[id_or_all || ::Rails.env]
       end
 
       def create(settings = {})
-        raise 'database servers should be set up in database.yml' unless ::Rails.env.test?
+        raise "database servers should be set up in database.yml" unless ::Rails.env.test?
 
         id = settings[:id]
         unless id
@@ -38,7 +39,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 +50,48 @@ module Switchman
         servers[rand(servers.length)]
       end
 
+      def guard_servers
+        all.each { |db| db.guard! if db.config[:prefer_secondary] }
+      end
+
+      def regions
+        @regions ||= all.filter_map(&:region).uniq.sort
+      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
         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('/')
+            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
+          return @database_servers if @database_servers.empty?
+
+          Shard.send(:configure_connects_to)
         end
         @database_servers
       end
@@ -102,8 +120,9 @@ module Switchman
       self.class.send(:database_servers).delete(id) if id
       Shard.sharded_models.each do |klass|
         self.class.all_roles.each do |role|
-          klass.connection_handler.remove_connection_pool(klass.connection_specification_name, role: role,
-                                                                                               shard: id.to_sym)
+          klass.connection_handler.remove_connection_pool(klass.connection_specification_name,
+                                                          role: role,
+                                                          shard: id.to_sym)
         end
       end
     end
@@ -129,32 +148,57 @@ module Switchman
         end
     end
 
-    def guard_rail_environment
-      @guard_rail_environment || ::GuardRail.environment
+    def region
+      config[:region]
+    end
+
+    # @param region [String, Array<String>] the region(s) to check against
+    # @return true if the database server doesn't have a region, or it
+    #   matches the specified region
+    def in_region?(region)
+      !self.region || (region.is_a?(Array) ? region.include?(self.region) : self.region == region)
+    end
+
+    # @return true if the database server doesn't have a region, Switchman is
+    #   not configured with a region, or the database server's region matches
+    #   Switchman's current region
+    def in_current_region?
+      unless instance_variable_defined?(:@in_current_region)
+        @in_current_region = !region ||
+                             !Switchman.region ||
+                             region == Switchman.region
+      end
+      @in_current_region
     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
       if id == ::Rails.env
-        Shard.where('database_server_id IS NULL OR database_server_id=?', id)
+        Shard.where("database_server_id IS NULL OR database_server_id=?", id)
       else
         Shard.where(database_server_id: id)
       end
@@ -179,7 +223,7 @@ module Switchman
       end
 
       id ||= begin
-        id_seq = Shard.connection.quote(Shard.connection.quote_table_name('switchman_shards_id_seq'))
+        id_seq = Shard.connection.quote(Shard.connection.quote_table_name("switchman_shards_id_seq"))
         next_id = Shard.connection.select_value("SELECT nextval(#{id_seq})")
         next_id.to_i
       end
@@ -196,29 +240,32 @@ module Switchman
                                 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)}")
+            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'
+              raise "This schema already exists; cannot overwrite"
             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
+          if config[:adapter] == "postgresql"
+            old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor {}
           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
+                if ::Rails.version < "7.1"
+                  ::ActiveRecord::Base.connection.migration_context.migrate
+                else
+                  ::ActiveRecord::MigrationContext.new(::ActiveRecord::Migrator.migrations_paths).migrate
+                end
               end
-              reset_column_information
+
               ::ActiveRecord::Base.descendants.reject do |m|
                 m <= UnshardedRecord || !m.table_exists?
               end.each(&:define_attribute_methods)
@@ -232,7 +279,6 @@ module Switchman
       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
@@ -257,18 +303,18 @@ module Switchman
     end
 
     def primary_shard
-      unless instance_variable_defined?(:@primary_shard)
-        # if sharding isn't fully set up yet, we may not be able to query the shards table
-        @primary_shard = Shard.default if Shard.default.database_server == self
-        @primary_shard ||= shards.where(name: nil).first
-      end
-      @primary_shard
-    end
+      return nil unless primary_shard_id
 
-    private
+      Shard.lookup(primary_shard_id)
+    end
 
-    def reset_column_information
-      ::ActiveRecord::Base.descendants.reject { |m| m <= UnshardedRecord }.each(&:reset_column_information)
+    def primary_shard_id
+      unless instance_variable_defined?(:@primary_shard_id)
+        # if sharding isn't fully set up yet, we may not be able to query the shards table
+        @primary_shard_id = Shard.default.id if Shard.default.database_server == self
+        @primary_shard_id ||= shards.where(name: nil).first&.id
+      end
+      @primary_shard_id
     end
   end
 end

data/lib/switchman/default_shard.rb

@@ -1,13 +1,12 @@
 # frozen_string_literal: true
 
-require 'switchman/database_server'
-
 module Switchman
   class DefaultShard
     def id
-      'default'
+      "default"
     end
-    alias cache_key id
+    alias_method :cache_key, :id
+
     def activate(*_classes)
       yield
     end
@@ -59,8 +58,18 @@ module Switchman
       self
     end
 
+    def region; end
+
+    def in_region?(_region)
+      true
+    end
+
+    def in_current_region?
+      true
+    end
+
     def _dump(_depth)
-      ''
+      ""
     end
 
     def self._load(_str)