switchman 3.0.1 → 4.2.5

data/lib/switchman/active_record/base.rb

@@ -6,16 +6,14 @@ module Switchman
       module ClassMethods
         delegate :shard, to: :all
 
-        def find_ids_in_ranges(opts = {}, &block)
+        def find_ids_in_ranges(opts = {}, &)
           opts.reverse_merge!(loose: true)
-          all.find_ids_in_ranges(opts, &block)
+          all.find_ids_in_ranges(opts, &)
         end
 
         def sharded_model
           self.abstract_class = true
 
-          return if self == UnshardedRecord
-
           Shard.send(:add_sharded_model, self)
         end
 
@@ -27,20 +25,12 @@ module Switchman
         def transaction(**)
           if self != ::ActiveRecord::Base && current_scope
             current_scope.activate do
-              db = Shard.current(connection_classes).database_server
-              if ::GuardRail.environment == db.guard_rail_environment
-                super
-              else
-                db.unguard { super }
-              end
-            end
-          else
-            db = Shard.current(connection_classes).database_server
-            if ::GuardRail.environment == db.guard_rail_environment
-              super
-            else
+              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
 
@@ -63,60 +53,206 @@ module Switchman
         end
 
         def clear_query_caches_for_current_thread
-          ::ActiveRecord::Base.connection_handler.connection_pool_list.each do |pool|
-            pool.connection(switch_shard: false).clear_query_cache if pool.active_connection?
+          pools = ::ActiveRecord::Base.connection_handler.connection_pool_list(:all)
+          pools.each do |pool|
+            if ::Rails.version < "7.2"
+              pool.connection(switch_shard: false).clear_query_cache if pool.active_connection?
+            elsif pool.active_connection?
+              pool.lease_connection(switch_shard: false).clear_query_cache
+            end
           end
         end
 
+        def role_overriden?(shard_id)
+          current_role(target_shard: shard_id) != current_role(without_overrides: true)
+        end
+
+        def establish_connection(config_or_env = nil)
+          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
+
+          config_or_env ||= if current_shard == ::Rails.env.to_sym && current_role == :primary
+                              :primary
+                            else
+                              :"#{current_shard}/#{current_role}"
+                            end
+
+          super
+        end
+
+        def connected_to_stack
+          has_own_stack = ::ActiveSupport::IsolatedExecutionState.key?(:active_record_connected_to_stack)
+
+          ret = super
+          return ret if has_own_stack
+
+          DatabaseServer.guard_servers
+          ret
+        end
+
+        # significant change: Allow per-shard roles
+        def current_role(without_overrides: false, target_shard: current_shard)
+          return super() if without_overrides
+
+          sharded_role = nil
+          connected_to_stack.reverse_each do |hash|
+            shard_role = hash.dig(:shard_roles, target_shard)
+            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
+
+          super()
+        end
+
         # significant change: _don't_ check if klasses.include?(Base)
         # i.e. other sharded models don't inherit the current shard of Base
         def current_shard
           connected_to_stack.reverse_each do |hash|
-            return hash[:shard] if hash[:shard] && hash[:klasses].include?(connection_classes)
+            return hash[:shard] if hash[:shard] && hash[:klasses].include?(connection_class_for_self)
           end
 
           default_shard
         end
+
+        def current_switchman_shard
+          connected_to_stack.reverse_each do |hash|
+            if hash[:switchman_shard] && hash[:klasses].include?(connection_class_for_self)
+              return hash[:switchman_shard]
+            end
+          end
+
+          Shard.default
+        end
       end
 
-      def self.included(klass)
+      def self.prepended(klass)
         klass.singleton_class.prepend(ClassMethods)
-        klass.set_callback(:initialize, :before) do
-          @shard ||= if self.class.sharded_primary_key?
-                       Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
-                     else
-                       Shard.current(self.class.connection_classes)
-                     end
+        klass.singleton_class.prepend(Switchman::ActiveRecord::Relation::InsertUpsertAll) if ::Rails.version < "7.2"
+        klass.scope :non_shadow, lambda { |key = primary_key|
+                                   where(key => (QueryMethods::NonTransposingValue.new(0)...
+                                                 QueryMethods::NonTransposingValue.new(Shard::IDS_PER_SHARD)))
+                                 }
+        klass.scope :shadow, lambda { |key = primary_key|
+                               where(key => QueryMethods::NonTransposingValue.new(Shard::IDS_PER_SHARD)..)
+                             }
+      end
+
+      def _run_initialize_callbacks
+        @shard ||= if self.class.sharded_primary_key?
+                     Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_class_for_self))
+                   else
+                     Shard.current(self.class.connection_class_for_self)
+                   end
+
+        @loaded_from_shard ||= Shard.current(self.class.connection_class_for_self)
+        if shadow_record? && !Switchman.config[:writable_shadow_records]
+          @readonly = true
+          @readonly_from_shadow ||= true
+        end
+        super
+      end
+
+      def readonly!
+        @readonly_from_shadow = false
+        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 canonical?
+        !shadow_record?
+      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_all([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_classes) || 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
-        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
+        super
       end
 
       def save!(*, **)
+        raise Errors::ManuallyCreatedShadowRecordError if creating_shadow_record?
+
         @shard_set_in_stone = true
-        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
+        super
       end
 
       def destroy
-        shard.activate(self.class.connection_classes) { super }
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def clone
@@ -128,14 +264,21 @@ module Switchman
         result
       end
 
-      def transaction(**kwargs, &block)
-        shard.activate(self.class.connection_classes) do
-          self.class.transaction(**kwargs, &block)
+      def transaction(...)
+        shard.activate(self.class.connection_class_for_self) do
+          self.class.transaction(...)
+        end
+      end
+
+      def with_transaction_returning_status
+        shard.activate(self.class.connection_class_for_self) do
+          db = Shard.current(self.class.connection_class_for_self).database_server
+          db.unguard { super }
         end
       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
@@ -149,42 +292,53 @@ module Switchman
         copy
       end
 
-      def quoted_id
-        return super unless self.class.sharded_primary_key?
-
-        # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
-        self.class.connection.quote(id)
+      def update_columns(*)
+        db = shard.database_server
+        db.unguard { super }
       end
 
-      def update_columns(*)
-        db = Shard.current(self.class.connection_classes).database_server
-        if ::GuardRail.environment == db.guard_rail_environment
-          super
+      def id_for_database
+        if self.class.sharded_primary_key?
+          # 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
-          db.unguard { super }
+          super
         end
       end
 
       protected
 
-      # see also AttributeMethods#connection_classes_code_for_reflection
-      def connection_classes_for_reflection(reflection)
+      # see also AttributeMethods#connection_class_for_self_code_for_reflection
+      def connection_class_for_self_for_reflection(reflection)
         if reflection
           if reflection.options[:polymorphic]
             begin
-              read_attribute(reflection.foreign_type)&.constantize&.connection_classes
+              read_attribute(reflection.foreign_type)&.constantize&.connection_class_for_self || ::ActiveRecord::Base
             rescue NameError
               # in case someone is abusing foreign_type to not point to an actual class
               ::ActiveRecord::Base
             end
           else
             # otherwise we can just return a symbol for the statically known type of the association
-            reflection.klass.connection_classes
+            reflection.klass.connection_class_for_self
           end
         else
-          connection_classes
+          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

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module Calculations
       def pluck(*column_names)
-        target_shard = Shard.current(klass.connection_classes)
+        target_shard = Shard.current(klass.connection_class_for_self)
         shard_count = 0
         result = activate do |relation, shard|
           shard_count += 1
@@ -28,19 +28,19 @@ 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|
+          result = activate(count: operation == "count") do |relation|
             relation.call_super(:execute_simple_calculation, Calculations, operation, column_name, distinct)
           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
@@ -50,22 +50,22 @@ module Switchman
 
       def calculate_simple_average(column_name, distinct)
         # See activerecord#execute_simple_calculation
-        relation = reorder(nil)
+        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
@@ -83,19 +83,19 @@ module Switchman
           if opts[:association]
             key_ids     = calculated_data.collect { |row| row[opts[:group_aliases].first] }
             key_records = opts[:association].klass.base_class.where(id: key_ids)
-            key_records = key_records.map { |r| [Shard.relative_id_for(r, shard, target_shard), r] }.to_h
+            key_records = key_records.to_h { |r| [Shard.relative_id_for(r, shard, target_shard), r] }
           end
 
           calculated_data.map do |row|
             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)
                 row[aliaz] = key_records[Shard.relative_id_for(row[aliaz], shard, target_shard)]
-              elsif group_column_name && @klass.sharded_column?(group_column_name)
+              elsif group_column_name && klass.sharded_column?(group_column_name)
                 row[aliaz] = Shard.relative_id_for(row[aliaz], shard, target_shard)
               end
             end
@@ -106,53 +106,95 @@ module Switchman
         compact_grouped_calculation_rows(rows, opts)
       end
 
+      def ids
+        return super unless klass.sharded_primary_key?
+
+        if loaded?
+          result = records.map do |record|
+            Shard.relative_id_for(record._read_attribute(primary_key),
+                                  record.shard,
+                                  Shard.current(klass.connection_class_for_self))
+          end
+          return @async ? Promise::Complete.new(result) : result
+        end
+
+        if has_include?(primary_key)
+          relation = apply_join_dependency.group(primary_key)
+          return relation.ids
+        end
+
+        columns = arel_columns([primary_key])
+        base_shard = Shard.current(klass.connection_class_for_self)
+        activate do |r|
+          relation = r.spawn
+          relation.select_values = columns
+
+          result = if relation.where_clause.contradiction?
+                     ::ActiveRecord::Result.empty
+                   else
+                     skip_query_cache_if_necessary do
+                       klass.connection.select_all(relation, "#{klass.name} Ids", async: @async)
+                     end
+                   end
+
+          result.then do |res|
+            type_cast_pluck_values(res, columns).map { |id| Shard.relative_id_for(id, Shard.current, base_shard) }
+          end
+        end
+      end
+
       private
 
       def type_cast_calculated_value_switchman(value, column_name, operation)
-        type_cast_calculated_value(value, operation) do |val|
-          column = aggregate_column(column_name)
-          type ||= column.try(:type_caster) ||
-                   lookup_cast_type_from_join_dependencies(column_name.to_s) || Type.default_value
-          type.deserialize(val)
-        end
+        column = aggregate_column(column_name)
+        type ||= column.try(:type_caster) ||
+                 lookup_cast_type_from_join_dependencies(column_name.to_s) || ::ActiveRecord::Type.default_value
+        type_cast_calculated_value(value, operation, type)
       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)
-        opts = { operation: operation, column_name: column_name, distinct: distinct }
+        opts = { operation:, column_name:, distinct: }
+
+        column_alias_tracker = ::ActiveRecord::Calculations::ColumnAliasTracker.new(connection)
 
-        opts[:aggregate_alias] = aggregate_alias_for(operation, column_name)
+        opts[:aggregate_alias] = aggregate_alias_for(operation, column_name, column_alias_tracker)
         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
         end
 
-        # to_s is because Rails 5 returns a string but Rails 6 returns a symbol.
-        group_aliases = group_fields.map { |field| column_alias_for(field.downcase.to_s).to_s }
+        group_aliases = group_fields.map do |field|
+          field = connection.visitor.compile(field) if ::Arel.arel_node?(field)
+          column_alias_tracker.alias_for(field.to_s.downcase)
+        end
         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,
-                    group_fields: group_fields)
+        opts.merge!(association:,
+                    associated:,
+                    group_aliases:,
+                    group_columns:,
+                    group_fields:)
 
         opts
       end
 
-      def aggregate_alias_for(operation, column_name)
-        if operation == 'count' && column_name == :all
-          'count_all'
-        elsif operation == 'average'
-          'average'
+      def aggregate_alias_for(operation, column_name, column_alias_tracker)
+        if operation == "count" && column_name == :all
+          "count_all"
+        elsif operation == "average"
+          "average"
         else
-          column_alias_for("#{operation} #{column_name}")
+          column_alias_tracker.alias_for("#{operation} #{column_name}")
         end
       end
 
@@ -166,13 +208,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)
@@ -198,22 +241,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
@@ -221,7 +264,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_handler.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module ConnectionHandler
+      def resolve_pool_config(config, connection_name, role, shard)
+        ret = super
+        # Make *all* pool configs use the same schema reflection
+        ret.schema_reflection = ConnectionHandler.global_schema_reflection
+        ret
+      end
+
+      def self.global_schema_reflection
+        @global_schema_reflection ||= ::ActiveRecord::ConnectionAdapters::SchemaReflection.new(nil)
+      end
+    end
+  end
+end