switchman 3.5.13 → 3.5.14

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c0ab0712ea0711e482d32ddf5ce6ec740b399b6dc8074ea86eaaa092c13bbdd1
-  data.tar.gz: a3e43929c88727da33b2954966f7cf357bd2b8793eed33563f57908839e2d97d
+  metadata.gz: 243c4f20483016069c4c18efd47a64a8b2078596f254c1152f4f2a70555054f0
+  data.tar.gz: 332d8e5a2d5e119d1599e259f27fa1ba25e9f3139cc7504bd159aebcb730e724
 SHA512:
-  metadata.gz: 67a005a7e156773576077ade7723e1500d291e782cdfb4af9662cef6ba2a217b2e0743aaea8a5e8bffc1777295baf465f9d7678f6d8fa6c681215b1ad9c5817a
-  data.tar.gz: 7e72531b558a3fc0d447d63efeec37b5a0480ebe873ff2e76fb6334927aee10d5342a92b0be8336bd324da57c2df6ff154b5efb2d88b5cc0e9e7b16f60bc6c78
+  metadata.gz: 53604b1eedb8bd5ca453d617eb7e4c83ad5080bad4ccb2dff0cced4be2baaeccbe139d7b5936face689e2ccbd94cac338d39f7ecf2ace8a25f61f81fe47ed0c8
+  data.tar.gz: 5f9b8523159c6eeec8275838a825c9c911ad136ef33d32d362e80ce63008f130e7789f1863b851d9527c68591bcbcfa6aa95183f2b6cca5eef7e55096194abb4

data/lib/switchman/active_record/base.rb

@@ -57,7 +57,12 @@ module Switchman
         end
 
         def clear_query_caches_for_current_thread
-          ::ActiveRecord::Base.connection_handler.connection_pool_list(:all).each do |pool|
+          pools = if ::Rails.version < "7.1"
+                    ::ActiveRecord::Base.connection_handler.connection_pool_list
+                  else
+                    ::ActiveRecord::Base.connection_handler.connection_pool_list(:all)
+                  end
+          pools.each do |pool|
             pool.connection(switch_shard: false).clear_query_cache if pool.active_connection?
           end
         end

data/lib/switchman/database_server.rb

@@ -10,6 +10,10 @@ module Switchman
       attr_accessor :creating_new_shard
       attr_reader :all_roles
 
+      include Enumerable
+
+      delegate :each, to: :all
+
       def all
         database_servers.values
       end
@@ -50,6 +54,10 @@ module Switchman
         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)
@@ -140,6 +148,29 @@ module Switchman
         end
     end
 
+    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)

data/lib/switchman/default_shard.rb

@@ -6,6 +6,7 @@ module Switchman
       "default"
     end
     alias_method :cache_key, :id
+
     def activate(*_classes)
       yield
     end
@@ -57,6 +58,16 @@ module Switchman
       self
     end
 
+    def region; end
+
+    def in_region?(_region)
+      true
+    end
+
+    def in_current_region?
+      true
+    end
+
     def _dump(_depth)
       ""
     end

data/lib/switchman/r_spec_helper.rb

@@ -121,7 +121,7 @@ module Switchman
         next if @@sharding_failed
 
         # clean up after specs
-        DatabaseServer.all.each do |ds|
+        DatabaseServer.each do |ds|
           if ds.fake? && ds != @shard2.database_server
             ds.shards.delete_all unless use_transactional_tests
             ds.destroy

data/lib/switchman/shard.rb

@@ -13,6 +13,47 @@ module Switchman
 
     scope :primary, -> { where(name: nil).order(:database_server_id, :id).distinct_on(:database_server_id) }
 
+    scope :in_region, (lambda do |region, include_regionless: true|
+      next in_current_region if region.nil?
+
+      dbs_by_region = DatabaseServer.group_by(&:region)
+      db_count_in_this_region = dbs_by_region[region]&.length.to_i
+      db_count_in_this_region += dbs_by_region[nil]&.length.to_i if include_regionless
+      non_existent_database_servers = Shard.send(:non_existent_database_servers)
+      db_count_in_other_regions = DatabaseServer.all.length -
+                                  db_count_in_this_region +
+                                  non_existent_database_servers.length
+
+      dbs_in_this_region = dbs_by_region[region]&.map(&:id) || []
+      dbs_in_this_region += dbs_by_region[nil]&.map(&:id) || [] if include_regionless
+
+      if db_count_in_this_region <= db_count_in_other_regions
+        if dbs_in_this_region.include?(Shard.default.database_server.id)
+          where("database_server_id IN (?) OR database_server_id IS NULL", dbs_in_this_region)
+        else
+          where(database_server_id: dbs_in_this_region)
+        end
+      elsif db_count_in_other_regions.zero?
+        all
+      else
+        dbs_not_in_this_region = DatabaseServer.map(&:id) - dbs_in_this_region + non_existent_database_servers
+        if dbs_in_this_region.include?(Shard.default.database_server.id)
+          where("database_server_id NOT IN (?) OR database_server_id IS NULL", dbs_not_in_this_region)
+        else
+          where.not(database_server_id: dbs_not_in_this_region)
+        end
+      end
+    end)
+
+    scope :in_current_region, (lambda do |include_regionless: true|
+      # sharding isn't set up? maybe we're in tests, or a somehow degraded environment
+      # either way there's only one shard, and we always want to see it
+      return [default] unless default.is_a?(Switchman::Shard)
+      return all if !Switchman.region || DatabaseServer.none?(&:region)
+
+      in_region(Switchman.region, include_regionless: include_regionless)
+    end)
+
     class << self
       def sharded_models
         @sharded_models ||= [::ActiveRecord::Base, UnshardedRecord].freeze
@@ -395,7 +436,7 @@ module Switchman
       end
 
       def configure_connects_to
-        full_connects_to_hash = DatabaseServer.all.to_h { |db| [db.id.to_sym, db.connects_to_hash] }
+        full_connects_to_hash = DatabaseServer.to_h { |db| [db.id.to_sym, db.connects_to_hash] }
         sharded_models.each do |klass|
           connects_to_hash = full_connects_to_hash.deep_dup
           if klass == UnshardedRecord
@@ -484,8 +525,17 @@ module Switchman
         argv[0] = File.basename(argv[0])
         argv.shelljoin
       end
+
+      # @return [Array<String>] the list of database servers that are in the
+      #   config, but don't have any shards on them
+      def non_existent_database_servers
+        @non_existent_database_servers ||=
+          Shard.distinct.pluck(:database_server_id).compact - DatabaseServer.all.map(&:id)
+      end
     end
 
+    delegate :region, :in_region?, :in_current_region?, to: :database_server
+
     def name
       unless instance_variable_defined?(:@name)
         # protect against re-entrancy

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = "3.5.13"
+  VERSION = "3.5.14"
 end

data/lib/switchman.rb

@@ -20,27 +20,31 @@ loader.setup
 module Switchman
   Deprecation = ::ActiveSupport::Deprecation.new("4.0", "Switchman")
 
-  def self.config
-    # TODO: load from yaml
-    @config ||= {}
-  end
+  class << self
+    attr_writer :cache
 
-  def self.cache
-    (@cache.respond_to?(:call) ? @cache.call : @cache) || ::Rails.cache
-  end
+    def config
+      # TODO: load from yaml
+      @config ||= {}
+    end
 
-  def self.cache=(cache)
-    @cache = cache
-  end
+    def cache
+      (@cache.respond_to?(:call) ? @cache.call : @cache) || ::Rails.cache
+    end
+
+    def region
+      config[:region]
+    end
 
-  def self.foreign_key_check(name, type, limit: nil)
-    return unless name.to_s.end_with?("_id") && type.to_s == "integer" && limit.to_i < 8
+    def foreign_key_check(name, type, limit: nil)
+      return unless name.to_s.end_with?("_id") && type.to_s == "integer" && limit.to_i < 8
 
-    puts <<~TEXT.squish
-      WARNING: All foreign keys need to be 8-byte integers.
-      #{name} looks like a foreign key.
-      If so, please add the option: `:limit => 8`
-    TEXT
+      puts <<~TEXT.squish
+        WARNING: All foreign keys need to be 8-byte integers.
+        #{name} looks like a foreign key.
+        If so, please add the option: `:limit => 8`
+      TEXT
+    end
   end
 
   class OrderOnMultiShardQuery < RuntimeError; end

data/lib/tasks/switchman.rake

@@ -35,7 +35,7 @@ module Switchman
 
         servers = servers.filter_map { |server| DatabaseServer.find(server) }
         if open
-          open_servers = DatabaseServer.all.select { |server| server.config[:open] }
+          open_servers = DatabaseServer.select { |server| server.config[:open] }
           servers.concat(open_servers)
           servers << DatabaseServer.find(nil) if open_servers.empty?
           servers.uniq!
@@ -43,6 +43,19 @@ module Switchman
         servers = DatabaseServer.all - servers if negative
       end
 
+      ENV["REGION"]&.split(",")&.each do |region|
+        method = :select!
+        if region[0] == "-"
+          method = :reject!
+          region = region[1..]
+        end
+        if region == "self"
+          servers.send(method, &:in_current_region?)
+        else
+          servers.send(method) { |server| server.in_region?(region) }
+        end
+      end
+
       servers = filter_database_servers_chain.call(servers)
 
       scope = base_scope.order(::Arel.sql("database_server_id IS NOT NULL, database_server_id, id"))

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.5.13
+  version: 3.5.14
 platform: ruby
 authors:
 - Cody Cutrer
@@ -10,7 +10,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2023-10-05 00:00:00.000000000 Z
+date: 2023-10-12 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord