switchman 3.1.0 → 3.5.0

data/lib/switchman/active_record/base.rb

@@ -22,16 +22,20 @@ module Switchman
           @integral_id
         end
 
-        def transaction(**)
-          if self != ::ActiveRecord::Base && current_scope
-            current_scope.activate do
-              db = Shard.current(connection_class_for_self).database_server
-              db.unguard { super }
+        %w[transaction insert_all upsert_all].each do |method|
+          class_eval <<-RUBY, __FILE__, __LINE__ + 1
+            def #{method}(*, **)
+              if self != ::ActiveRecord::Base && current_scope
+                current_scope.activate do
+                  db = Shard.current(connection_class_for_self).database_server
+                  db.unguard { super }
+                end
+              else
+                db = Shard.current(connection_class_for_self).database_server
+                db.unguard { super }
+              end
             end
-          else
-            db = Shard.current(connection_class_for_self).database_server
-            db.unguard { super }
-          end
+          RUBY
         end
 
         def reset_column_information
@@ -63,7 +67,10 @@ module Switchman
         end
 
         def establish_connection(config_or_env = nil)
-          raise ArgumentError, 'establish connection cannot be used on the non-current shard/role' if config_or_env.is_a?(Symbol) && config_or_env != ::Rails.env.to_sym
+          if config_or_env.is_a?(Symbol) && config_or_env != ::Rails.env.to_sym
+            raise ArgumentError,
+                  "establish connection cannot be used on the non-current shard/role"
+          end
 
           # Ensure we don't randomly surprise change the connection parms associated with a shard/role
           config_or_env = nil if config_or_env == ::Rails.env.to_sym
@@ -78,9 +85,15 @@ module Switchman
         end
 
         def connected_to_stack
-          return super if ::Rails.version < '7.0' ? Thread.current.thread_variable?(:ar_connected_to_stack) : ::ActiveSupport::IsolatedExecutionState.key?(:active_record_connected_to_stack)
+          has_own_stack = if ::Rails.version < "7.0"
+                            Thread.current.thread_variable?(:ar_connected_to_stack)
+                          else
+                            ::ActiveSupport::IsolatedExecutionState.key?(:active_record_connected_to_stack)
+                          end
 
           ret = super
+          return ret if has_own_stack
+
           DatabaseServer.guard_servers
           ret
         end
@@ -92,10 +105,13 @@ module Switchman
           sharded_role = nil
           connected_to_stack.reverse_each do |hash|
             shard_role = hash.dig(:shard_roles, target_shard)
-            if shard_role && (hash[:klasses].include?(::ActiveRecord::Base) || hash[:klasses].include?(connection_class_for_self))
-              sharded_role = shard_role
-              break
+            unless shard_role &&
+                   (hash[:klasses].include?(::ActiveRecord::Base) || hash[:klasses].include?(connection_class_for_self))
+              next
             end
+
+            sharded_role = shard_role
+            break
           end
           # Allow a shard-specific role to be reverted to regular inheritance
           return sharded_role if sharded_role && sharded_role != :_switchman_inherit
@@ -115,13 +131,15 @@ module Switchman
 
         def current_switchman_shard
           connected_to_stack.reverse_each do |hash|
-            return hash[:switchman_shard] if hash[:switchman_shard] && hash[:klasses].include?(connection_class_for_self)
+            if hash[:switchman_shard] && hash[:klasses].include?(connection_class_for_self)
+              return hash[:switchman_shard]
+            end
           end
 
           Shard.default
         end
 
-        if ::Rails.version < '7.0'
+        if ::Rails.version < "7.0"
           def connection_class_for_self
             connection_classes
           end
@@ -138,30 +156,86 @@ module Switchman
                    else
                      Shard.current(self.class.connection_class_for_self)
                    end
+
+        @loaded_from_shard ||= Shard.current(self.class.connection_class_for_self)
+        readonly! if shadow_record? && !Switchman.config[:writable_shadow_records]
         super
       end
 
+      def shadow_record?
+        pkey = self[self.class.primary_key]
+        return false unless self.class.sharded_column?(self.class.primary_key) && pkey
+
+        pkey > Shard::IDS_PER_SHARD
+      end
+
+      def save_shadow_record(new_attrs: attributes, target_shard: Shard.current)
+        return if target_shard == shard
+
+        shadow_attrs = {}
+        new_attrs.each do |attr, value|
+          shadow_attrs[attr] = if self.class.sharded_column?(attr)
+                                 Shard.relative_id_for(value, shard, target_shard)
+                               else
+                                 value
+                               end
+        end
+        target_shard.activate do
+          self.class.upsert(shadow_attrs, unique_by: self.class.primary_key)
+        end
+      end
+
+      def destroy_shadow_records(target_shards: [Shard.current])
+        raise Errors::ShadowRecordError, "Cannot be called on a shadow record." if shadow_record?
+
+        unless self.class.sharded_column?(self.class.primary_key)
+          raise Errors::MethodUnsupportedForUnshardedTableError,
+                "Cannot be called on a record belonging to an unsharded table."
+        end
+
+        Array(target_shards).each do |target_shard|
+          next if target_shard == shard
+
+          target_shard.activate { self.class.where("id = ?", global_id).delete_all }
+        end
+      end
+
+      # Returns "the shard that this record was actually loaded from" , as
+      # opposed to "the shard this record belongs on", which might be
+      # different if this is a shadow record.
+      def loaded_from_shard
+        @loaded_from_shard || shard
+      end
+
       def shard
-        @shard || Shard.current(self.class.connection_class_for_self) || Shard.default
+        @shard || fallback_shard
       end
 
       def shard=(new_shard)
         raise ::ActiveRecord::ReadOnlyRecord if !new_record? || @shard_set_in_stone
 
-        return if shard == new_shard
+        if shard == new_shard
+          @loaded_from_shard = new_shard
+          return
+        end
 
         attributes.each do |attr, value|
           self[attr] = Shard.relative_id_for(value, shard, new_shard) if self.class.sharded_column?(attr)
         end
+        @loaded_from_shard = new_shard
         @shard = new_shard
       end
 
       def save(*, **)
+        raise Errors::ManuallyCreatedShadowRecordError if creating_shadow_record?
+
         @shard_set_in_stone = true
         super
       end
 
       def save!(*, **)
+        raise Errors::ManuallyCreatedShadowRecordError if creating_shadow_record?
+
         @shard_set_in_stone = true
         super
       end
@@ -193,7 +267,7 @@ module Switchman
       end
 
       def hash
-        self.class.sharded_primary_key? ? self.class.hash ^ global_id.hash : super
+        self.class.sharded_primary_key? ? [self.class, global_id].hash : super
       end
 
       def to_param
@@ -214,8 +288,10 @@ module Switchman
 
       def id_for_database
         if self.class.sharded_primary_key?
-          # It's an int, so so it's safe to just return it without passing it through anything else
-          # In theory we should do `@attributes[@primary_key].type.serialize(id)`, but that seems to have surprising side-effects
+          # It's an int, so it's safe to just return it without passing it
+          # through anything else. In theory we should do
+          # `@attributes[@primary_key].type.serialize(id)`, but that seems to
+          # have surprising side-effects
           id
         else
           super
@@ -242,6 +318,16 @@ module Switchman
           self.class.connection_class_for_self
         end
       end
+
+      private
+
+      def fallback_shard
+        Shard.current(self.class.connection_class_for_self) || Shard.default
+      end
+
+      def creating_shadow_record?
+        new_record? && shadow_record?
+      end
     end
   end
 end

data/lib/switchman/active_record/calculations.rb

@@ -28,7 +28,7 @@ module Switchman
 
       def execute_simple_calculation(operation, column_name, distinct)
         operation = operation.to_s.downcase
-        if operation == 'average'
+        if operation == "average"
           result = calculate_simple_average(column_name, distinct)
         else
           result = activate do |relation|
@@ -36,11 +36,11 @@ module Switchman
           end
           if result.is_a?(Array)
             case operation
-            when 'count', 'sum'
+            when "count", "sum"
               result = result.sum
-            when 'minimum'
+            when "minimum"
               result = result.min
-            when 'maximum'
+            when "maximum"
               result = result.max
             end
           end
@@ -52,20 +52,20 @@ module Switchman
         # See activerecord#execute_simple_calculation
         relation = except(:order)
         column = aggregate_column(column_name)
-        relation.select_values = [operation_over_aggregate_column(column, 'average', distinct).as('average'),
-                                  operation_over_aggregate_column(column, 'count', distinct).as('count')]
+        relation.select_values = [operation_over_aggregate_column(column, "average", distinct).as("average"),
+                                  operation_over_aggregate_column(column, "count", distinct).as("count")]
 
         initial_results = relation.activate { |rel| klass.connection.select_all(rel) }
         if initial_results.is_a?(Array)
           initial_results.each do |r|
-            r['average'] = type_cast_calculated_value_switchman(r['average'], column_name, 'average')
-            r['count'] = type_cast_calculated_value_switchman(r['count'], column_name, 'count')
+            r["average"] = type_cast_calculated_value_switchman(r["average"], column_name, "average")
+            r["count"] = type_cast_calculated_value_switchman(r["count"], column_name, "count")
           end
-          result = initial_results.map { |r| r['average'] * r['count'] }.sum / initial_results.map do |r|
-                                                                                 r['count']
-                                                                               end.sum
+          result = initial_results.sum { |r| r["average"] * r["count"] } / initial_results.sum do |r|
+                                                                             r["count"]
+                                                                           end
         else
-          result = type_cast_calculated_value_switchman(initial_results.first['average'], column_name, 'average')
+          result = type_cast_calculated_value_switchman(initial_results.first["average"], column_name, "average")
         end
         result
       end
@@ -90,7 +90,7 @@ module Switchman
             row[opts[:aggregate_alias]] = type_cast_calculated_value_switchman(
               row[opts[:aggregate_alias]], column_name, opts[:operation]
             )
-            row['count'] = row['count'].to_i if opts[:operation] == 'average'
+            row["count"] = row["count"].to_i if opts[:operation] == "average"
 
             opts[:group_columns].each do |aliaz, _type, group_column_name|
               if opts[:associated] && (aliaz == opts[:group_aliases].first)
@@ -109,7 +109,7 @@ module Switchman
       private
 
       def type_cast_calculated_value_switchman(value, column_name, operation)
-        if ::Rails.version < '7.0'
+        if ::Rails.version < "7.0"
           type_cast_calculated_value(value, operation) do |val|
             column = aggregate_column(column_name)
             type ||= column.try(:type_caster) ||
@@ -125,7 +125,7 @@ module Switchman
       end
 
       def column_name_for(field)
-        field.respond_to?(:name) ? field.name.to_s : field.to_s.split('.').last
+        field.respond_to?(:name) ? field.name.to_s : field.to_s.split(".").last
       end
 
       def grouped_calculation_options(operation, column_name, distinct)
@@ -135,7 +135,8 @@ module Switchman
         group_attrs = group_values
         if group_attrs.first.respond_to?(:to_sym)
           association  = klass.reflect_on_association(group_attrs.first.to_sym)
-          associated   = group_attrs.size == 1 && association && association.macro == :belongs_to # only count belongs_to associations
+          # only count belongs_to associations
+          associated   = group_attrs.size == 1 && association && association.macro == :belongs_to
           group_fields = Array(associated ? association.foreign_key : group_attrs)
         else
           group_fields = group_attrs
@@ -146,18 +147,20 @@ module Switchman
         group_columns = group_aliases.zip(group_fields).map do |aliaz, field|
           [aliaz, type_for(field), column_name_for(field)]
         end
-        opts.merge!(association: association, associated: associated,
-                    group_aliases: group_aliases, group_columns: group_columns,
+        opts.merge!(association: association,
+                    associated: associated,
+                    group_aliases: group_aliases,
+                    group_columns: group_columns,
                     group_fields: group_fields)
 
         opts
       end
 
       def aggregate_alias_for(operation, column_name)
-        if operation == 'count' && column_name == :all
-          'count_all'
-        elsif operation == 'average'
-          'average'
+        if operation == "count" && column_name == :all
+          "count_all"
+        elsif operation == "average"
+          "average"
         else
           column_alias_for("#{operation} #{column_name}")
         end
@@ -173,13 +176,14 @@ module Switchman
             opts[:distinct]
           ).as(opts[:aggregate_alias])
         ]
-        if opts[:operation] == 'average'
+        if opts[:operation] == "average"
           # include count in average so we can recalculate the average
           # across all shards if needed
           select_values << operation_over_aggregate_column(
             aggregate_column(opts[:column_name]),
-            'count', opts[:distinct]
-          ).as('count')
+            "count",
+            opts[:distinct]
+          ).as("count")
         end
 
         haves = having_clause.send(:predicates)
@@ -205,22 +209,22 @@ module Switchman
           key = key.first if key.size == 1
           value = row[opts[:aggregate_alias]]
 
-          if opts[:operation] == 'average'
+          if opts[:operation] == "average"
             if result.key?(key)
               old_value, old_count = result[key]
-              new_count = old_count + row['count']
-              new_value = ((old_value * old_count) + (value * row['count'])) / new_count
+              new_count = old_count + row["count"]
+              new_value = ((old_value * old_count) + (value * row["count"])) / new_count
               result[key] = [new_value, new_count]
             else
-              result[key] = [value, row['count']]
+              result[key] = [value, row["count"]]
             end
           elsif result.key?(key)
             case opts[:operation]
-            when 'count', 'sum'
+            when "count", "sum"
               result[key] += value
-            when 'minimum'
+            when "minimum"
               result[key] = value if value < result[key]
-            when 'maximum'
+            when "maximum"
               result[key] = value if value > result[key]
             end
           else
@@ -228,7 +232,7 @@ module Switchman
           end
         end
 
-        result.transform_values!(&:first) if opts[:operation] == 'average'
+        result.transform_values!(&:first) if opts[:operation] == "average"
 
         result
       end

data/lib/switchman/active_record/connection_pool.rb

@@ -3,6 +3,23 @@
 module Switchman
   module ActiveRecord
     module ConnectionPool
+      def get_schema_cache(connection)
+        self.schema_cache ||= SharedSchemaCache.get_schema_cache(connection)
+        self.schema_cache.connection = connection
+
+        self.schema_cache
+      end
+
+      # rubocop:disable Naming/AccessorMethodName override method
+      def set_schema_cache(cache)
+        schema_cache = get_schema_cache(cache.connection)
+
+        cache.instance_variables.each do |x|
+          schema_cache.instance_variable_set(x, cache.instance_variable_get(x))
+        end
+      end
+      # rubocop:enable Naming/AccessorMethodName override method
+
       def default_schema
         connection unless @schemas
         # default shard will not switch databases immediately, so it won't be set yet
@@ -31,7 +48,9 @@ module Switchman
       end
 
       def switch_database(conn)
-        @schemas = conn.current_schemas if !@schemas && conn.adapter_name == 'PostgreSQL' && !current_shard.database_server.config[:shard_name]
+        if !@schemas && conn.adapter_name == "PostgreSQL" && !current_shard.database_server.config[:shard_name]
+          @schemas = conn.current_schemas
+        end
 
         conn.shard = current_shard
       end
@@ -39,7 +58,7 @@ module Switchman
       private
 
       def current_shard
-        ::Rails.version < '7.0' ? connection_klass.current_switchman_shard : connection_class.current_switchman_shard
+        (::Rails.version < "7.0") ? connection_klass.current_switchman_shard : connection_class.current_switchman_shard
       end
 
       def tls_key

data/lib/switchman/active_record/database_configurations.rb

@@ -10,9 +10,9 @@ module Switchman
       def configs_for(include_replicas: false, name: nil, **)
         res = super
         if name && !include_replicas
-          return nil unless name.end_with?('primary')
+          return nil unless name.end_with?("primary")
         elsif !include_replicas
-          return res.select { |config| config.name.end_with?('primary') }
+          return res.select { |config| config.name.end_with?("primary") }
         end
         res
       end
@@ -29,19 +29,24 @@ module Switchman
         db_configs = configs.flat_map do |env_name, config|
           # It would be nice to do the auto-fallback that we want here, but we haven't
           # actually done that for years (or maybe ever) and it will be a big lift to get working
-          roles = config.keys.select { |k| config[k].is_a?(Hash) || (config[k].is_a?(Array) && config[k].all? { |ck| ck.is_a?(Hash) }) }
+          roles = config.keys.select do |k|
+            config[k].is_a?(Hash) || (config[k].is_a?(Array) && config[k].all?(Hash))
+          end
           base_config = config.except(*roles)
 
           name = "#{env_name}/primary"
-          name = 'primary' if env_name == default_env
+          name = "primary" if env_name == default_env
           base_db = build_db_config_from_raw_config(env_name, name, base_config)
           [base_db] + roles.map do |role|
-            build_db_config_from_raw_config(env_name, "#{env_name}/#{role}",
-                                            base_config.merge(config[role].is_a?(Array) ? config[role].first : config[role]))
+            build_db_config_from_raw_config(
+              env_name,
+              "#{env_name}/#{role}",
+              base_config.merge(config[role].is_a?(Array) ? config[role].first : config[role])
+            )
           end
         end
 
-        db_configs << environment_url_config(default_env, 'primary', {}) unless db_configs.find(&:for_current_env?)
+        db_configs << environment_url_config(default_env, "primary", {}) unless db_configs.find(&:for_current_env?)
 
         merge_db_environment_variables(default_env, db_configs.compact)
       end

data/lib/switchman/active_record/finder_methods.rb

@@ -49,7 +49,7 @@ module Switchman
         relation = apply_join_dependency(eager_loading: false)
         return false if ::ActiveRecord::NullRelation === relation
 
-        relation = relation.except(:select, :order).select('1 AS one').limit(1)
+        relation = relation.except(:select, :order).select("1 AS one").limit(1)
 
         case conditions
         when Array, Hash

data/lib/switchman/active_record/log_subscriber.rb

@@ -14,14 +14,14 @@ module Switchman
 
         name  = "#{payload[:name]} (#{event.duration.round(1)}ms)"
         name  = "CACHE #{name}" if payload[:cached]
-        sql   = payload[:sql].squeeze(' ')
+        sql   = payload[:sql].squeeze(" ")
         binds = nil
         shard = payload[:shard]
         shard = "  [#{shard[:database_server_id]}:#{shard[:id]} #{shard[:env]}]" if shard
 
         unless (payload[:binds] || []).empty?
           casted_params = type_casted_binds(payload[:type_casted_binds])
-          binds = '  ' + payload[:binds].zip(casted_params).map do |attr, value|
+          binds = "  " + payload[:binds].zip(casted_params).map do |attr, value|
             render_bind(attr, value)
           end.inspect
         end

data/lib/switchman/active_record/migration.rb

@@ -14,19 +14,23 @@ module Switchman
 
       def connection
         conn = super
-        ::ActiveRecord::Base.connection_pool.switch_database(conn) if conn.shard != ::ActiveRecord::Base.current_switchman_shard
+        if conn.shard != ::ActiveRecord::Base.current_switchman_shard
+          ::ActiveRecord::Base.connection_pool.switch_database(conn)
+        end
         conn
       end
     end
 
     module Migrator
-      # significant change: just return MIGRATOR_SALT directly
-      # especially if you're going through pgbouncer, the database
-      # name you're accessing may not be consistent. it is NOT allowed
-      # to run migrations against multiple shards in the same database
-      # concurrently
+      # significant change: use the shard name instead of the database name
+      # in the lock id. Especially if you're going through pgbouncer, the
+      # database name you're accessing may not be consistent
       def generate_migrator_advisory_lock_id
-        ::ActiveRecord::Migrator::MIGRATOR_SALT
+        db_name_hash = Zlib.crc32(Shard.current.name)
+        shard_name_hash = ::ActiveRecord::Migrator::MIGRATOR_SALT * db_name_hash
+        # Store in internalmetadata to allow other tools to be able to lock out migrations
+        ::ActiveRecord::InternalMetadata[:migrator_advisory_lock_id] = shard_name_hash
+        shard_name_hash
       end
 
       # significant change: strip out prefer_secondary from config
@@ -42,13 +46,36 @@ module Switchman
     end
 
     module MigrationContext
+      def migrate(...)
+        connection = ::ActiveRecord::Base.connection
+        connection_pool = ::ActiveRecord::Base.connection_pool
+        previous_schema_cache = connection_pool.get_schema_cache(connection)
+        temporary_schema_cache = ::ActiveRecord::ConnectionAdapters::SchemaCache.new(connection)
+
+        reset_column_information
+        connection_pool.set_schema_cache(temporary_schema_cache)
+
+        begin
+          super(...)
+        ensure
+          connection_pool.set_schema_cache(previous_schema_cache)
+          reset_column_information
+        end
+      end
+
       def migrations
         return @migrations if instance_variable_defined?(:@migrations)
 
         migrations_cache = Thread.current[:migrations_cache] ||= {}
-        key = Digest::MD5.hexdigest(migration_files.sort.join(','))
+        key = Digest::MD5.hexdigest(migration_files.sort.join(","))
         @migrations = migrations_cache[key] ||= super
       end
+
+      private
+
+      def reset_column_information
+        ::ActiveRecord::Base.descendants.reject { |m| m <= UnshardedRecord }.each(&:reset_column_information)
+      end
     end
   end
 end

data/lib/switchman/active_record/persistence.rb

@@ -16,6 +16,14 @@ module Switchman
         db = shard.database_server
         db.unguard { super }
       end
+
+      def reload(*)
+        res = super
+        # When a shadow record is reloaded the real record is returned. So
+        # we need to ensure the loaded_from_shard is set correctly after a reload.
+        @loaded_from_shard = @shard
+        res
+      end
     end
   end
 end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -5,9 +5,9 @@ module Switchman
     module PostgreSQLAdapter
       # copy/paste; use quote_local_table_name
       def create_database(name, options = {})
-        options = { encoding: 'utf8' }.merge!(options.symbolize_keys)
+        options = { encoding: "utf8" }.merge!(options.symbolize_keys)
 
-        option_string = options.sum('') do |key, value|
+        option_string = options.sum("") do |key, value|
           case key
           when :owner
             " OWNER = \"#{value}\""
@@ -24,7 +24,7 @@ module Switchman
           when :connection_limit
             " CONNECTION LIMIT = #{value}"
           else
-            ''
+            ""
           end
         end
 
@@ -37,7 +37,7 @@ module Switchman
       end
 
       def current_schemas
-        select_values('SELECT * FROM unnest(current_schemas(false))')
+        select_values("SELECT * FROM unnest(current_schemas(false))")
       end
 
       def extract_schema_qualified_name(string)
@@ -49,13 +49,13 @@ module Switchman
       # significant change: use the shard name if no explicit schema
       def quoted_scope(name = nil, type: nil)
         schema, name = extract_schema_qualified_name(name)
-        type = \
+        type =
           case type # rubocop:disable Style/HashLikeCase
-          when 'BASE TABLE'
+          when "BASE TABLE"
             "'r','p'"
-          when 'VIEW'
+          when "VIEW"
             "'v','m'"
-          when 'FOREIGN TABLE'
+          when "FOREIGN TABLE"
             "'f'"
           end
         scope = {}
@@ -67,7 +67,8 @@ module Switchman
 
       def foreign_keys(table_name)
         super.each do |fk|
-          to_table_qualified_name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(fk.to_table)
+          to_table_qualified_name =
+            ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(fk.to_table)
           fk.to_table = to_table_qualified_name.identifier if to_table_qualified_name.schema == shard.name
         end
       end
@@ -101,7 +102,7 @@ module Switchman
 
       def add_index_options(_table_name, _column_name, **)
         index, algorithm, if_not_exists = super
-        algorithm = nil if DatabaseServer.creating_new_shard && algorithm == 'CONCURRENTLY'
+        algorithm = nil if DatabaseServer.creating_new_shard && algorithm == "CONCURRENTLY"
         [index, algorithm, if_not_exists]
       end
 

data/lib/switchman/active_record/query_cache.rb

@@ -20,7 +20,7 @@ module Switchman
                 type_casted_binds: -> { type_casted_binds(binds) }
               }
               ::ActiveSupport::Notifications.instrument(
-                'sql.active_record',
+                "sql.active_record",
                 args
               )
               query_cache[sql][binds]