switchman 3.4.2 → 3.6.7

data/lib/switchman/active_record/finder_methods.rb

@@ -42,26 +42,56 @@ module Switchman
         primary_shard.activate { super }
       end
 
-      def exists?(conditions = :none)
-        conditions = conditions.id if ::ActiveRecord::Base === conditions
-        return false unless conditions
+      if ::Rails.version < "7.1"
+        def exists?(conditions = :none)
+          conditions = conditions.id if ::ActiveRecord::Base === conditions
+          return false unless conditions
 
-        relation = apply_join_dependency(eager_loading: false)
-        return false if ::ActiveRecord::NullRelation === relation
+          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
-          relation = relation.where(conditions)
-        else
-          relation = relation.where(table[primary_key].eq(conditions)) if conditions != :none
+          case conditions
+          when Array, Hash
+            relation = relation.where(conditions)
+          else
+            relation = relation.where(table[primary_key].eq(conditions)) if conditions != :none
+          end
+
+          relation.activate do |shard_rel|
+            return true if connection.select_value(shard_rel.arel, "#{name} Exists")
+          end
+          false
         end
+      else
+        def exists?(conditions = :none)
+          return false if @none
+
+          if Base === conditions
+            raise ArgumentError, <<-TEXT.squish
+              You are passing an instance of ActiveRecord::Base to `exists?`.
+              Please pass the id of the object by calling `.id`.
+            TEXT
+          end
 
-        relation.activate do |shard_rel|
-          return true if connection.select_value(shard_rel.arel, "#{name} Exists")
+          return false if !conditions || limit_value == 0 # rubocop:disable Style/NumericPredicate
+
+          if eager_loading?
+            relation = apply_join_dependency(eager_loading: false)
+            return relation.exists?(conditions)
+          end
+
+          relation = construct_relation_for_exists(conditions)
+          return false if relation.where_clause.contradiction?
+
+          relation.activate do |shard_rel|
+            return true if skip_query_cache_if_necessary do
+                             connection.select_rows(shard_rel.arel, "#{name} Exists?").size == 1
+                           end
+          end
+          false
         end
-        false
       end
     end
   end

data/lib/switchman/active_record/log_subscriber.rb

@@ -5,8 +5,10 @@ module Switchman
     module LogSubscriber
       # sadly, have to completely replace this
       def sql(event)
-        self.class.runtime += event.duration
-        return unless logger.debug?
+        if ::Rails.version < "7.1"
+          self.class.runtime += event.duration
+          return unless logger.debug?
+        end
 
         payload = event.payload
 
@@ -14,20 +16,24 @@ 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
 
         name = colorize_payload_name(name, payload[:name])
-        sql  = color(sql, sql_color(sql), true)
+        sql  = if ::Rails.version < "7.1"
+                 color(sql, sql_color(sql), true)
+               else
+                 color(sql, sql_color(sql), bold: true)
+               end
 
         debug "  #{name}  #{sql}#{binds}#{shard}"
       end

data/lib/switchman/active_record/migration.rb

@@ -14,7 +14,9 @@ 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
@@ -27,7 +29,11 @@ module Switchman
         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
+        if ::Rails.version < "7.1"
+          ::ActiveRecord::InternalMetadata[:migrator_advisory_lock_id] = shard_name_hash
+        else
+          ::ActiveRecord::InternalMetadata.new(connection)[:migrator_advisory_lock_id] = shard_name_hash
+        end
         shard_name_hash
       end
 
@@ -46,17 +52,30 @@ module Switchman
     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)
+        schema_cache_holder = ::ActiveRecord::Base.connection_pool
+        schema_cache_holder = schema_cache_holder.schema_reflection if ::Rails.version >= "7.1"
+        previous_schema_cache = if ::Rails.version < "7.1"
+                                  schema_cache_holder.get_schema_cache(connection)
+                                else
+                                  schema_cache_holder.instance_variable_get(:@cache)
+                                end
 
-        reset_column_information
-        connection_pool.set_schema_cache(temporary_schema_cache)
+        if ::Rails.version < "7.1"
+          temporary_schema_cache = ::ActiveRecord::ConnectionAdapters::SchemaCache.new(connection)
+
+          reset_column_information
+          schema_cache_holder.set_schema_cache(temporary_schema_cache)
+        else
+          schema_cache_holder.instance_variable_get(:@cache)
+
+          reset_column_information
+          schema_cache_holder.clear!
+        end
 
         begin
           super(...)
         ensure
-          connection_pool.set_schema_cache(previous_schema_cache)
+          schema_cache_holder.set_schema_cache(previous_schema_cache)
           reset_column_information
         end
       end
@@ -65,7 +84,7 @@ module Switchman
         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
 

data/lib/switchman/active_record/pending_migration_connection.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module PendingMigrationConnection
+      module ClassMethods
+        def current_role
+          ::ActiveRecord::Base.current_role
+        end
+
+        def current_switchman_shard
+          ::ActiveRecord::Base.current_switchman_shard
+        end
+      end
+    end
+  end
+end

data/lib/switchman/active_record/persistence.rb

@@ -5,10 +5,12 @@ module Switchman
     module Persistence
       # touch reads the id attribute directly, so it's not relative to the current shard
       def touch(*, **)
+        writable_shadow_record_warning
         shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def update_columns(*)
+        writable_shadow_record_warning
         shard.activate(self.class.connection_class_for_self) { super }
       end
 
@@ -17,13 +19,33 @@ module Switchman
         db.unguard { super }
       end
 
+      def destroy
+        writable_shadow_record_warning
+        super
+      end
+
+      def create_or_update(**, &block)
+        writable_shadow_record_warning
+        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
+        if @readonly_from_shadow
+          @readonly_from_shadow = false
+          @readonly = false
+        end
         res
       end
+
+      def writable_shadow_record_warning
+        return unless shadow_record? && Switchman.config[:writable_shadow_records]
+
+        Switchman::Deprecation.warn("writing to shadow records is not supported")
+      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/predicate_builder.rb

@@ -11,11 +11,11 @@ module Switchman
         end
       end
 
-      module AssociationQueryValue
+      module PolymorphicArrayValue
         def convert_to_id(value)
           case value
           when ::ActiveRecord::Base
-            value.id
+            value.send(primary_key(value))
           else
             super
           end

data/lib/switchman/active_record/query_cache.rb

@@ -5,28 +5,57 @@ module Switchman
     module QueryCache
       private
 
-      def cache_sql(sql, name, binds)
-        # have to include the shard id in the cache key because of switching dbs on the same connection
-        sql = "#{shard.id}::#{sql}"
-        @lock.synchronize do
-          result =
-            if query_cache[sql].key?(binds)
-              args = {
-                sql: sql,
-                binds: binds,
-                name: name,
-                connection_id: object_id,
-                cached: true,
-                type_casted_binds: -> { type_casted_binds(binds) }
-              }
-              ::ActiveSupport::Notifications.instrument(
-                'sql.active_record',
-                args
-              )
-              query_cache[sql][binds]
+      if ::Rails.version < "7.1"
+        def cache_sql(sql, name, binds)
+          # have to include the shard id in the cache key because of switching dbs on the same connection
+          sql = "#{shard.id}::#{sql}"
+          @lock.synchronize do
+            result =
+              if query_cache[sql].key?(binds)
+                args = {
+                  sql: sql,
+                  binds: binds,
+                  name: name,
+                  connection_id: object_id,
+                  cached: true,
+                  type_casted_binds: -> { type_casted_binds(binds) }
+                }
+                ::ActiveSupport::Notifications.instrument(
+                  "sql.active_record",
+                  args
+                )
+                query_cache[sql][binds]
+              else
+                query_cache[sql][binds] = yield
+              end
+            result.dup
+          end
+        end
+      else
+        def cache_sql(sql, name, binds)
+          # have to include the shard id in the cache key because of switching dbs on the same connection
+          sql = "#{shard.id}::#{sql}"
+          key = binds.empty? ? sql : [sql, binds]
+          result = nil
+          hit = false
+
+          @lock.synchronize do
+            if (result = @query_cache.delete(key))
+              hit = true
+              @query_cache[key] = result
             else
-              query_cache[sql][binds] = yield
+              result = @query_cache[key] = yield
+              @query_cache.shift if @query_cache_max_size && @query_cache.size > @query_cache_max_size
             end
+          end
+
+          if hit
+            ::ActiveSupport::Notifications.instrument(
+              "sql.active_record",
+              cache_notification_info(sql, name, binds)
+            )
+          end
+
           result.dup
         end
       end

data/lib/switchman/active_record/query_methods.rb

@@ -3,6 +3,19 @@
 module Switchman
   module ActiveRecord
     module QueryMethods
+      # Use this class to prevent a value from getting transposed across shards
+      class NonTransposingValue < SimpleDelegator
+        def class
+          __getobj__.class
+        end
+
+        def is_a?(other)
+          return true if other == NonTransposingValue
+
+          __getobj__.is_a?(other)
+        end
+      end
+
       # shard_value is one of:
       #   A shard
       #   An array or relation of shards
@@ -84,6 +97,14 @@ module Switchman
         super(other.shard(primary_shard))
       end
 
+      # use a temp variable so that the new where clause is built before self.where_clause is read,
+      # since build_where_clause might mutate self.where_clause
+      def where!(opts, *rest)
+        new_clause = build_where_clause(opts, rest)
+        self.where_clause += new_clause
+        self
+      end
+
       protected
 
       def remove_nonlocal_primary_keys!
@@ -93,7 +114,7 @@ module Switchman
             predicate.left.relation.klass == klass &&
             (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::HomogeneousIn))
 
-          value.is_a?(Integer) && value > Shard::IDS_PER_SHARD ? [] : value
+          (value.is_a?(Integer) && value > Shard::IDS_PER_SHARD) ? [] : value
         end
         self
       end
@@ -104,7 +125,9 @@ module Switchman
         return unless klass.integral_id?
 
         primary_key = predicates.detect do |predicate|
-          (predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)) &&
+          (predicate.is_a?(::Arel::Nodes::Equality) ||
+           predicate.is_a?(::Arel::Nodes::In) ||
+           predicate.is_a?(::Arel::Nodes::HomogeneousIn)) &&
             predicate.left.is_a?(::Arel::Attributes::Attribute) &&
             predicate.left.relation.is_a?(::Arel::Table) && predicate.left.relation.klass == klass &&
             klass.primary_key == predicate.left.name
@@ -167,19 +190,19 @@ module Switchman
       end
 
       def sharded_foreign_key?(relation, column)
-        models_for_table(relation.table_name).any? { |m| m.sharded_column?(column) }
+        models_for_table(relation.name).any? { |m| m.sharded_column?(column) }
       end
 
       def sharded_primary_key?(relation, column)
         column = column.to_s
-        return column == 'id' if relation.klass == ::ActiveRecord::Base
+        return column == "id" if relation.klass == ::ActiveRecord::Base
 
         relation.klass.primary_key == column && relation.klass.integral_id?
       end
 
       def source_shard_for_foreign_key(relation, column)
         reflection = nil
-        models_for_table(relation.table_name).each do |model|
+        models_for_table(relation.name).each do |model|
           reflection = model.send(:reflection_for_integer_attribute, column)
           break if reflection
         end
@@ -199,12 +222,11 @@ module Switchman
 
         case opts
         when String, Array
-          values = Hash === rest.first ? rest.first.values : rest
+          values = (Hash === rest.first) ? rest.first.values : rest
 
-          values.grep(ActiveRecord::Relation) do |rel|
-            # serialize subqueries against the same shard as the outer query is currently
-            # targeted to run against
-            rel.shard!(primary_shard) if rel.shard_source_value == :implicit && rel.primary_shard != primary_shard
+          if shard_source_value != :explicit && values.grep(ActiveRecord::Relation).first
+            raise "Sub-queries are not allowed as simple substitutions; " \
+                  "please build your relation with more structured methods so that Switchman is able to introspect it."
           end
 
           super
@@ -254,9 +276,10 @@ module Switchman
         send(clause_setter, new_clause)
       end
 
-      def each_transposable_predicate(predicates = nil, &block)
+      def each_transposable_predicate(predicates, &block)
         each_predicate(predicates) do |predicate|
-          if predicate.is_a?(::Arel::Nodes::Grouping)
+          case predicate
+          when ::Arel::Nodes::Grouping
             next predicate unless predicate.expr.is_a?(::Arel::Nodes::Or)
 
             or_expr = predicate.expr
@@ -267,6 +290,40 @@ module Switchman
             next predicate if new_left == old_left && new_right == old_right
 
             next predicate.class.new predicate.expr.class.new(new_left, new_right)
+          when ::Arel::Nodes::SelectStatement
+            new_cores = predicate.cores.map do |core|
+              next core unless core.is_a?(::Arel::Nodes::SelectCore) # just in case something weird is going on
+
+              new_wheres = each_transposable_predicate(core.wheres, &block)
+              new_havings = each_transposable_predicate(core.havings, &block)
+
+              next core if core.wheres == new_wheres && core.havings == new_havings
+
+              new_core = core.clone
+              new_core.wheres = new_wheres
+              new_core.havings = new_havings
+              new_core
+            end
+
+            next predicate if predicate.cores == new_cores
+
+            new_node = predicate.clone
+            new_node.instance_variable_set(:@cores, new_cores)
+            next new_node
+          when ::Arel::Nodes::Not
+            old_value = predicate.expr
+            new_value = each_transposable_predicate([old_value], &block).first
+
+            next predicate if old_value == new_value
+
+            next predicate.class.new(new_value)
+          when ::Arel::Nodes::Exists
+            old_value = predicate.expressions
+            new_value = each_transposable_predicate([old_value], &block).first
+
+            next predicate if old_value == new_value
+
+            next predicate.class.new(new_value)
           end
 
           next predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
@@ -279,54 +336,56 @@ module Switchman
         end
       end
 
-      def each_transposable_predicate_value(predicates = nil)
+      def each_transposable_predicate_value(predicates = nil, &block)
         each_transposable_predicate(predicates) do |predicate, relation, column, type|
-          each_transposable_predicate_value_cb(predicate) do |value|
+          each_transposable_predicate_value_cb(predicate, block) do |value|
             yield(value, predicate, relation, column, type)
           end
         end
       end
 
-      def each_transposable_predicate_value_cb(node, &block)
+      def each_transposable_predicate_value_cb(node, original_block, &block)
         case node
         when Array
-          node.map { |val| each_transposable_predicate_value_cb(val, &block).presence }.compact
+          node.filter_map { |val| each_transposable_predicate_value_cb(val, original_block, &block).presence }
         when ::ActiveModel::Attribute
           old_value = node.value_before_type_cast
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(node.name, new_value, node.type)
+          (old_value == new_value) ? node : node.class.new(node.name, new_value, node.type)
         when ::Arel::Nodes::And
           old_value = node.children
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(new_value)
+          (old_value == new_value) ? node : node.class.new(new_value)
         when ::Arel::Nodes::BindParam
           old_value = node.value
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(new_value)
+          (old_value == new_value) ? node : node.class.new(new_value)
         when ::Arel::Nodes::Casted
           old_value = node.value
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(new_value, node.attribute)
+          (old_value == new_value) ? node : node.class.new(new_value, node.attribute)
         when ::Arel::Nodes::HomogeneousIn
           old_value = node.values
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
           # switch to a regular In, so that Relation::WhereClause#contradiction? knows about it
           if new_value.empty?
-            klass = node.type == :in ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
+            klass = (node.type == :in) ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
             klass.new(node.attribute, new_value)
           else
-            old_value == new_value ? node : node.class.new(new_value, node.attribute, node.type)
+            (old_value == new_value) ? node : node.class.new(new_value, node.attribute, node.type)
           end
         when ::Arel::Nodes::Binary
           old_value = node.right
-          new_value = each_transposable_predicate_value_cb(old_value, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
 
-          old_value == new_value ? node : node.class.new(node.left, new_value)
+          (old_value == new_value) ? node : node.class.new(node.left, new_value)
+        when ::Arel::Nodes::SelectStatement
+          each_transposable_predicate_value([node], &original_block).first
         else
           yield(node)
         end
@@ -343,6 +402,8 @@ module Switchman
               Shard.current(klass.connection_class_for_self)
             elsif type == :foreign
               source_shard_for_foreign_key(relation, column)
+            else
+              primary_shard
             end
 
           transpose_predicate_value(value, current_source_shard, target_shard, type)
@@ -350,7 +411,9 @@ module Switchman
       end
 
       def transpose_predicate_value(value, current_shard, target_shard, attribute_type)
-        if value.is_a?(::ActiveRecord::StatementCache::Substitute)
+        if value.is_a?(NonTransposingValue)
+          value
+        elsif value.is_a?(::ActiveRecord::StatementCache::Substitute)
           value.sharded = true # mark for transposition later
           value.primary = true if attribute_type == :primary
           value