switchman 3.0.1 → 4.2.5

data/lib/switchman/sharded_instrumenter.rb

@@ -3,7 +3,7 @@
 module Switchman
   class ShardedInstrumenter < ::SimpleDelegator
     def initialize(instrumenter, shard_host)
-      super instrumenter
+      super(instrumenter)
       @shard_host = shard_host
     end
 
@@ -13,13 +13,19 @@ module Switchman
       # when we might be doing a query while defining attribute methods,
       # so just avoid logging then
       if shard.is_a?(Shard) && Shard.instance_variable_get(:@attribute_methods_generated)
+        env = if ::Rails.version < "8.0"
+                @shard_host.pool.connection_class&.current_role
+              else
+                @shard_host.pool.connection_descriptor&.name&.constantize&.current_role
+              end
+
         payload[:shard] = {
           database_server_id: shard.database_server.id,
           id: shard.id,
-          env: shard.database_server.guard_rail_environment
+          env:
         }
       end
-      super name, payload
+      super
     end
   end
 end

data/lib/switchman/shared_schema_cache.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Switchman
+  class SharedSchemaCache
+    def self.get_schema_cache(connection)
+      @schema_cache ||= ::ActiveRecord::ConnectionAdapters::SchemaCache.new(connection)
+      @schema_cache.connection = connection
+      @schema_cache
+    end
+  end
+end

data/lib/switchman/standard_error.rb

@@ -3,20 +3,23 @@
 module Switchman
   module StandardError
     def initialize(*args)
-      # Shard.current can throw this when switchman isn't working right; if we try to
-      # do our stuff here, it'll cause a SystemStackError, which is a pain to deal with
-      if is_a?(::ActiveRecord::ConnectionNotEstablished)
-        super
-        return
-      end
+      super
+      # These seem to get themselves into a bad state if we try to lookup shards while processing
+      return if is_a?(IO::EAGAINWaitReadable)
 
-      if defined?(Shard)
-        @active_shards = Shard.sharded_models.map do |klass|
-          [klass, Shard.current(klass)]
-        end.compact.to_h
-      end
+      return if Thread.current[:switchman_error_handler]
 
-      super
+      begin
+        Thread.current[:switchman_error_handler] = true
+
+        @active_shards ||= Shard.active_shards
+      rescue
+        # intentionally empty - don't allow calculating the active_shards to prevent
+        # creating the StandardError for any reason. this prevents various random issues
+        # when a StandardError is created within a finalizer
+      ensure
+        Thread.current[:switchman_error_handler] = nil
+      end
     end
 
     def current_shard(klass = ::ActiveRecord::Base)

data/lib/switchman/test_helper.rb

@@ -19,7 +19,7 @@ module Switchman
         end
 
         server1 = Shard.default.database_server
-        server2 = DatabaseServer.create(Shard.default.database_server.config.merge(server2: true))
+        server2 = DatabaseServer.create(Shard.default.database_server.config.merge(server2: true, schema_dump: false))
 
         if server1 == Shard.default.database_server && server1.config[:shard1] && server1.config[:shard2]
           # look for the shards in the db already
@@ -63,9 +63,9 @@ module Switchman
 
       def find_existing_test_shard(server, name)
         if server == Shard.default.database_server
-          server.shards.where(name: name).first
+          server.shards.where(name:).first
         else
-          shard = Shard.where('database_server_id IS NOT NULL AND name=?', name).first
+          shard = Shard.where("database_server_id IS NOT NULL AND name=?", name).first
           # if somehow databases got created in a different order, change the shard to match
           shard.database_server = server if shard
           shard

data/{app/models → lib}/switchman/unsharded_record.rb

@@ -2,6 +2,6 @@
 
 module Switchman
   class UnshardedRecord < ::ActiveRecord::Base
-    sharded_model
+    self.abstract_class = true
   end
 end

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = '3.0.1'
+  VERSION = "4.2.5"
 end

data/lib/switchman.rb

@@ -1,20 +1,54 @@
 # frozen_string_literal: true
 
-require 'guard_rail'
-require 'switchman/open4'
-require 'switchman/engine'
+require "guard_rail"
+require "zeitwerk"
 
-module Switchman
-  def self.config
-    # TODO: load from yaml
-    @config ||= {}
+class SwitchmanInflector < Zeitwerk::GemInflector
+  def camelize(basename, abspath)
+    if basename =~ /\Apostgresql_(.*)/
+      "PostgreSQL" + super($1, abspath)
+    else
+      super
+    end
   end
+end
 
-  def self.cache
-    (@cache.respond_to?(:call) ? @cache.call : @cache) || ::Rails.cache
-  end
+loader = Zeitwerk::Loader.for_gem
+loader.inflector = SwitchmanInflector.new(__FILE__)
+loader.setup
+
+module Switchman
+  Deprecation = ::ActiveSupport::Deprecation.new("4.0", "Switchman")
+
+  class << self
+    attr_writer :cache
 
-  def self.cache=(cache)
-    @cache = cache
+    def config
+      # TODO: load from yaml
+      @config ||= {}
+    end
+
+    def cache
+      (@cache.respond_to?(:call) ? @cache.call : @cache) || ::Rails.cache
+    end
+
+    def region
+      config[:region]
+    end
+
+    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
+    end
   end
+
+  class OrderOnMultiShardQuery < RuntimeError; end
 end
+
+# Load the engine and everything associated at gem load time
+Switchman::Engine

data/lib/tasks/switchman.rake

@@ -2,28 +2,29 @@
 
 module Switchman
   module Rake
-    def self.filter_database_servers(&block)
-      chain = filter_database_servers_chain # use a local variable so that the current chain is closed over in the following lambda
-      @filter_database_servers_chain = ->(servers) { block.call(servers, chain) }
+    def self.filter_database_servers
+      # use a local variable so that the current chain is closed over in the following lambda
+      chain = filter_database_servers_chain
+      @filter_database_servers_chain = ->(servers) { yield(servers, chain) }
     end
 
     def self.scope(base_scope = Shard,
-                   database_server: ENV['DATABASE_SERVER'],
-                   shard: ENV['SHARD'])
+                   database_server: ENV.fetch("DATABASE_SERVER", nil),
+                   shard: ENV.fetch("SHARD", nil))
       servers = DatabaseServer.all
 
       if database_server
         servers = database_server
-        if servers.first == '-'
+        if servers.first == "-"
           negative = true
           servers = servers[1..]
         end
-        servers = servers.split(',')
-        open = servers.delete('open')
+        servers = servers.split(",")
+        open = servers.delete("open")
 
-        servers = servers.map { |server| DatabaseServer.find(server) }.compact
+        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!
@@ -31,13 +32,26 @@ 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'))
+      scope = base_scope.order(::Arel.sql("database_server_id IS NOT NULL, database_server_id, id"))
       if servers != DatabaseServer.all
-        conditions = ['database_server_id IN (?)', servers.map(&:id)]
-        conditions.first << ' OR database_server_id IS NULL' if servers.include?(Shard.default.database_server)
-        scope = scope.where(conditions)
+        database_server_ids = servers.map(&:id)
+        database_server_ids << nil if servers.include?(Shard.default.database_server)
+        scope = scope.where(database_server_id: database_server_ids)
       end
 
       scope = shard_scope(scope, shard) if shard
@@ -46,52 +60,93 @@ module Switchman
     end
 
     def self.options
-      { parallel: ENV['PARALLEL'].to_i, max_procs: ENV['MAX_PARALLEL_PROCS'] }
+      { exception: (ENV["FAIL_FAST"] == "0") ? :defer : :raise, parallel: ENV["PARALLEL"].to_i }
     end
 
     # classes - an array or proc, to activate as the current shard during the
-    # task. tasks which modify the schema may want to pass all categories in
-    # so that schema updates for non-default tables happen against all shards.
-    # this is handled automatically for the default migration tasks, below.
+    # task.
     def self.shardify_task(task_name, classes: [::ActiveRecord::Base])
+      log_format = ENV.fetch("LOG_FORMAT", nil)
       old_task = ::Rake::Task[task_name]
       old_actions = old_task.actions.dup
       old_task.actions.clear
 
       old_task.enhance do |*task_args|
         if ::Rails.env.test?
-          require 'switchman/test_helper'
+          require "switchman/test_helper"
           TestHelper.recreate_persistent_test_shards(dont_create: true)
         end
 
         ::GuardRail.activate(:deploy) do
           Shard.default.database_server.unguard do
             classes = classes.call if classes.respond_to?(:call)
-            Shard.with_each_shard(scope, classes, **options) do
+
+            # We don't want the shard status messages to be wrapped using a custom log transfomer
+            original_stderr = $stderr
+            original_stdout = $stdout
+            output = if log_format == "json"
+                       lambda { |msg|
+                         JSON.dump(shard: Shard.current.id,
+                                   database_server: Shard.current.database_server.id,
+                                   type: "log",
+                                   message: msg)
+                       }
+                     else
+                       nil
+                     end
+            Shard.with_each_shard(scope, classes, output:, **options) do
               shard = Shard.current
-              puts "#{shard.id}: #{shard.description}"
+
+              if log_format == "json"
+                original_stdout.puts JSON.dump(
+                  shard: shard.id,
+                  database_server: shard.database_server.id,
+                  type: "started"
+                )
+              else
+                original_stdout.puts "#{shard.id}: #{shard.description}"
+              end
 
               shard.database_server.unguard do
                 old_actions.each { |action| action.call(*task_args) }
               end
+
+              if log_format == "json"
+                original_stdout.puts JSON.dump(
+                  shard: shard.id,
+                  database_server: shard.database_server.id,
+                  type: "completed"
+                )
+              end
               nil
+            rescue => e
+              if log_format == "json"
+                original_stderr.puts JSON.dump(
+                  shard: shard.id,
+                  database_server: shard.database_server.id,
+                  type: "failed",
+                  message: e.full_message
+                )
+              end
+
+              raise
             end
           rescue => e
-            puts "Exception from #{e.current_shard.id}: #{e.current_shard.description}" if options[:parallel] != 0
+            if options[:parallel] != 0
+              warn "Exception from #{e.current_shard.id}: #{e.current_shard.description}:\n#{e.full_message}"
+            end
             raise
-
-            #::ActiveRecord::Base.configurations = old_configurations
           end
         end
       end
     end
 
     %w[db:migrate db:migrate:up db:migrate:down db:rollback].each do |task_name|
-      shardify_task(task_name, classes: -> { Shard.sharded_models })
+      shardify_task(task_name)
     end
 
     def self.shard_scope(scope, raw_shard_ids)
-      raw_shard_ids = raw_shard_ids.split(',')
+      raw_shard_ids = raw_shard_ids.split(",")
 
       shard_ids = []
       negative_shard_ids = []
@@ -101,13 +156,13 @@ module Switchman
 
       raw_shard_ids.each do |id|
         case id
-        when 'default'
+        when "default"
           shard_ids << Shard.default.id
-        when '-default'
+        when "-default"
           negative_shard_ids << Shard.default.id
-        when 'primary'
+        when "primary"
           shard_ids.concat(Shard.primary.pluck(:id))
-        when '-primary'
+        when "-primary"
           negative_shard_ids.concat(Shard.primary.pluck(:id))
         when /^(-?)(\d+)?\.\.(\.)?(\d+)?$/
           negative, start, open, finish = $1.present?, $2, $3.present?, $4
@@ -115,8 +170,8 @@ module Switchman
 
           range = []
           range << "id>=#{start}" if start
-          range << "id<#{'=' unless open}#{finish}" if finish
-          (negative ? negative_ranges : ranges) << "(#{range.join(' AND ')})"
+          range << "id<#{"=" unless open}#{finish}" if finish
+          (negative ? negative_ranges : ranges) << "(#{range.join(" AND ")})"
         when /^-(\d+)$/
           negative_shard_ids << $1.to_i
         when /^\d+$/
@@ -144,21 +199,21 @@ module Switchman
           select = []
           if index != 1
             subscope = subscope.offset(per_chunk * (index - 1))
-            select << 'MIN(id) AS min_id'
+            select << "MIN(id) AS min_id"
           end
           if index != denominator
             subscope = subscope.limit(per_chunk)
-            select << 'MAX(id) AS max_id'
+            select << "MAX(id) AS max_id"
           end
 
-          result = Shard.from(subscope).select(select.join(', ')).to_a.first
+          result = Shard.from(subscope).select(select.join(", ")).to_a.first
           range = case index
                   when 1
-                    "id<=#{result['max_id']}"
+                    "id<=#{result["max_id"]}"
                   when denominator
-                    "id>=#{result['min_id']}"
+                    "id>=#{result["min_id"]}"
                   else
-                    "(id>=#{result['min_id']} AND id<=#{result['max_id']})"
+                    "(id>=#{result["min_id"]} AND id<=#{result["max_id"]})"
                   end
 
           (numerator.negative? ? negative_ranges : ranges) << range
@@ -179,16 +234,16 @@ module Switchman
 
       conditions = []
       positive_queries = []
-      positive_queries << ranges.join(' OR ') unless ranges.empty?
+      positive_queries << ranges.join(" OR ") unless ranges.empty?
       unless shard_ids.empty?
-        positive_queries << 'id IN (?)'
+        positive_queries << "id IN (?)"
         conditions << shard_ids
       end
-      positive_query = positive_queries.join(' OR ')
+      positive_query = positive_queries.join(" OR ")
       scope = scope.where(positive_query, *conditions) unless positive_queries.empty?
 
-      scope = scope.where("NOT (#{negative_ranges.join(' OR')})") unless negative_ranges.empty?
-      scope = scope.where('id NOT IN (?)', negative_shard_ids) unless negative_shard_ids.empty?
+      scope = scope.where("NOT (#{negative_ranges.join(" OR")})") unless negative_ranges.empty?
+      scope = scope.where("id NOT IN (?)", negative_shard_ids) unless negative_shard_ids.empty?
       scope
     end
 
@@ -199,17 +254,9 @@ module Switchman
 
   module ActiveRecord
     module PostgreSQLDatabaseTasks
-      def structure_dump(filename, extra_flags = nil)
-        set_psql_env
-        args = ['-s', '-x', '-O', '-f', filename]
-        args.concat(Array(extra_flags)) if extra_flags
-        shard = Shard.current.name
-        serialized_search_path = shard
-        args << "--schema=#{Shellwords.escape(shard)}"
-
-        args << configuration['database']
-        run_cmd('pg_dump', args, 'dumping')
-        File.open(filename, 'a') { |f| f << "SET search_path TO #{serialized_search_path};\n\n" }
+      def structure_dump(...)
+        ::ActiveRecord.dump_schemas = Switchman::Shard.current.name
+        super
       end
     end
   end