switchman 4.0.0 → 4.2.0

data/lib/switchman/active_record/connection_pool.rb

@@ -3,29 +3,11 @@
 module Switchman
   module ActiveRecord
     module ConnectionPool
-      if ::Rails.version < "7.1"
-        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
-      end
-
       def default_schema
-        connection unless @schemas
+        connection_method = (::Rails.version < "7.2") ? :connection : :lease_connection
+        send(connection_method) unless @schemas
         # default shard will not switch databases immediately, so it won't be set yet
-        @schemas ||= connection.current_schemas
+        @schemas ||= send(connection_method).current_schemas
         @schemas.first
       end
 
@@ -43,8 +25,36 @@ module Switchman
         conn
       end
 
+      unless ::Rails.version < "7.2"
+        def active_connection(switch_shard: true)
+          conn = super()
+          return nil if conn.nil?
+          raise Errors::NonExistentShardError if current_shard.new_record?
+
+          switch_database(conn) if conn.shard != current_shard && switch_shard
+          conn
+        end
+
+        def lease_connection(switch_shard: true)
+          conn = super()
+          raise Errors::NonExistentShardError if current_shard.new_record?
+
+          switch_database(conn) if conn.shard != current_shard && switch_shard
+          conn
+        end
+
+        def with_connection(switch_shard: true, **kwargs)
+          super(**kwargs) do |conn|
+            raise Errors::NonExistentShardError if current_shard.new_record?
+
+            switch_database(conn) if conn.shard != current_shard && switch_shard
+            yield conn
+          end
+        end
+      end
+
       def release_connection(with_id = Thread.current)
-        super(with_id)
+        super
 
         flush
       end
@@ -60,7 +70,11 @@ module Switchman
       private
 
       def current_shard
-        connection_class.current_switchman_shard
+        if ::Rails.version < "8.0"
+          connection_class.current_switchman_shard
+        else
+          connection_descriptor.name.constantize.current_switchman_shard
+        end
       end
 
       def tls_key

data/lib/switchman/active_record/database_configurations.rb

@@ -7,26 +7,14 @@ module Switchman
       # since all should point to the same data, even if multiple are writable
       # (Picks 'primary' since it is guaranteed to exist and switchman handles activating
       # deploy through other means)
-      if ::Rails.version < "7.1"
-        def configs_for(include_replicas: false, name: nil, **)
-          res = super
-          if name && !include_replicas
-            return nil unless name.end_with?("primary")
-          elsif !include_replicas
-            return res.select { |config| config.name.end_with?("primary") }
-          end
-          res
-        end
-      else
-        def configs_for(include_hidden: false, name: nil, **)
-          res = super
-          if name && !include_hidden
-            return nil unless name.end_with?("primary")
-          elsif !include_hidden
-            return res.select { |config| config.name.end_with?("primary") }
-          end
-          res
+      def configs_for(include_hidden: false, name: nil, **)
+        res = super
+        if name && !include_hidden
+          return nil unless name.end_with?("primary")
+        elsif !include_hidden
+          return res.select { |config| config.name.end_with?("primary") }
         end
+        res
       end
 
       private

data/lib/switchman/active_record/finder_methods.rb

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module FinderMethods
       def find_one(id)
-        return super(id) unless klass.integral_id?
+        return super unless klass.integral_id?
 
         if shard_source_value != :implicit
           current_shard = Shard.current(klass.connection_class_for_self)
@@ -28,70 +28,46 @@ module Switchman
         if shard
           shard.activate { super(local_id) }
         else
-          super(id)
+          super
         end
       end
 
       def find_some_ordered(ids)
         current_shard = Shard.current(klass.connection_class_for_self)
         ids = ids.map { |id| Shard.relative_id_for(id, current_shard, current_shard) }
-        super(ids)
+        super
       end
 
       def find_or_instantiator_by_attributes(match, attributes, *args)
         primary_shard.activate { super }
       end
 
-      if ::Rails.version < "7.1"
-        def exists?(conditions = :none)
-          conditions = conditions.id if ::ActiveRecord::Base === conditions
-          return false unless conditions
+      def exists?(conditions = :none)
+        return false if @none
 
-          relation = apply_join_dependency(eager_loading: false)
-          return false if ::ActiveRecord::NullRelation === relation
-
-          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
-          end
-
-          relation.activate do |shard_rel|
-            return true if connection.select_value(shard_rel.arel, "#{name} Exists")
-          end
-          false
+        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
-      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
-
-          return false if !conditions || limit_value == 0 # rubocop:disable Style/NumericPredicate
+        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
+        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 = 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
+        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
-          false
         end
+        false
       end
     end
   end

data/lib/switchman/active_record/log_subscriber.rb

@@ -5,11 +5,6 @@ module Switchman
     module LogSubscriber
       # sadly, have to completely replace this
       def sql(event)
-        if ::Rails.version < "7.1"
-          self.class.runtime += event.duration
-          return unless logger.debug?
-        end
-
         payload = event.payload
 
         return if ::ActiveRecord::LogSubscriber::IGNORE_PAYLOAD_NAMES.include?(payload[:name])
@@ -29,11 +24,7 @@ module Switchman
         end
 
         name = colorize_payload_name(name, payload[:name])
-        sql  = if ::Rails.version < "7.1"
-                 color(sql, sql_color(sql), true)
-               else
-                 color(sql, sql_color(sql), bold: true)
-               end
+        sql  = color(sql, sql_color(sql), bold: true)
 
         debug "  #{name}  #{sql}#{binds}#{shard}"
       end

data/lib/switchman/active_record/migration.rb

@@ -29,11 +29,7 @@ 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
-        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
+        @internal_metadata[:migrator_advisory_lock_id] = shard_name_hash.to_s
         shard_name_hash
       end
 
@@ -51,31 +47,23 @@ module Switchman
 
     module MigrationContext
       def migrate(...)
-        connection = ::ActiveRecord::Base.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
-
-        if ::Rails.version < "7.1"
-          temporary_schema_cache = ::ActiveRecord::ConnectionAdapters::SchemaCache.new(connection)
+        schema_cache_holder = schema_cache_holder.schema_reflection
+        previous_schema_cache = schema_cache_holder.instance_variable_get(:@cache)
 
-          reset_column_information
-          schema_cache_holder.set_schema_cache(temporary_schema_cache)
-        else
-          schema_cache_holder.instance_variable_get(:@cache)
+        schema_cache_holder.instance_variable_get(:@cache)
 
-          reset_column_information
-          schema_cache_holder.clear!
-        end
+        reset_column_information
+        schema_cache_holder.clear!
 
         begin
-          super(...)
+          super
         ensure
-          schema_cache_holder.set_schema_cache(previous_schema_cache)
+          if ::Rails.version < "7.2"
+            schema_cache_holder.set_schema_cache(previous_schema_cache)
+          else
+            schema_cache_holder.instance_variable_set(:@cache, previous_schema_cache)
+          end
           reset_column_information
         end
       end

data/lib/switchman/active_record/persistence.rb

@@ -24,7 +24,7 @@ module Switchman
         super
       end
 
-      def create_or_update(**, &block)
+      def create_or_update(**, &)
         writable_shadow_record_warning
         super
       end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -73,23 +73,37 @@ module Switchman
         end
       end
 
+      module ClassMethods
+        def quote_local_table_name(name)
+          # postgres quotes tables and columns the same; just pass through
+          # (differs from quote_table_name_with_shard below by no logic to
+          # explicitly qualify the table)
+          quote_column_name(name)
+        end
+
+        def quote_table_name(name, shard: nil)
+          # This looks kind of weird at first glance, but older Rails versions do not actually import
+          # these methods as class methods.
+          shard = self.shard if shard.nil? && ::Rails.version < "7.2" && !@use_local_table_name
+
+          return quote_local_table_name(name) unless shard
+
+          name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
+          name.instance_variable_set(:@schema, shard.name) unless name.schema
+          name.quoted
+        end
+      end
+
       def quote_local_table_name(name)
-        # postgres quotes tables and columns the same; just pass through
-        # (differs from quote_table_name_with_shard below by no logic to
-        # explicitly qualify the table)
-        quote_column_name(name)
+        self.class.quote_local_table_name(name)
       end
 
       def quote_table_name(name)
-        return quote_local_table_name(name) if @use_local_table_name
-
-        name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
-        name.instance_variable_set(:@schema, shard.name) unless name.schema
-        name.quoted
+        self.class.quote_table_name(name, shard: @use_local_table_name ? nil : shard)
       end
 
-      def with_global_table_name(&block)
-        with_local_table_name(false, &block)
+      def with_global_table_name(&)
+        with_local_table_name(false, &)
       end
 
       def with_local_table_name(enable = true) # rubocop:disable Style/OptionalBooleanParameter
@@ -129,7 +143,7 @@ module Switchman
       end
 
       def columns(*)
-        with_local_table_name(false) { super }
+        with_global_table_name { super }
       end
     end
   end

data/lib/switchman/active_record/query_cache.rb

@@ -5,41 +5,15 @@ module Switchman
     module QueryCache
       private
 
-      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
+      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
+        @lock.synchronize do
+          if ::Rails.version < "7.2"
             if (result = @query_cache.delete(key))
               hit = true
               @query_cache[key] = result
@@ -47,17 +21,23 @@ module Switchman
               result = @query_cache[key] = yield
               @query_cache.shift if @query_cache_max_size && @query_cache.size > @query_cache_max_size
             end
+          else
+            hit = true
+            result = @query_cache.compute_if_absent(key) do
+              hit = false
+              yield
+            end
           end
+        end
 
-          if hit
-            ::ActiveSupport::Notifications.instrument(
-              "sql.active_record",
-              cache_notification_info(sql, name, binds)
-            )
-          end
-
-          result.dup
+        if hit
+          ::ActiveSupport::Notifications.instrument(
+            "sql.active_record",
+            cache_notification_info(sql, name, binds)
+          )
         end
+
+        result.dup
       end
     end
   end

data/lib/switchman/active_record/query_methods.rb

@@ -34,13 +34,29 @@ module Switchman
       end
 
       def shard_value=(value)
-        raise ::ActiveRecord::ImmutableRelation if @loaded
+        if @loaded
+          error_class = if ::Rails.version < "7.2"
+                          ::ActiveRecord::ImmutableRelation
+                        else
+                          ::ActiveRecord::UnmodifiableRelation
+                        end
+
+          raise error_class
+        end
 
         @values[:shard] = value
       end
 
       def shard_source_value=(value)
-        raise ::ActiveRecord::ImmutableRelation if @loaded
+        if @loaded
+          error_class = if ::Rails.version < "7.2"
+                          ::ActiveRecord::ImmutableRelation
+                        else
+                          ::ActiveRecord::UnmodifiableRelation
+                        end
+
+          raise error_class
+        end
 
         @values[:shard_source] = value
       end
@@ -107,6 +123,10 @@ module Switchman
 
       protected
 
+      def arel_columns(columns)
+        connection.with_local_table_name { super }
+      end
+
       def remove_nonlocal_primary_keys!
         each_transposable_predicate_value do |value, predicate, _relation, _column, type|
           next value unless
@@ -243,31 +263,49 @@ module Switchman
         end
       end
 
-      def arel_columns(columns)
-        connection.with_local_table_name { super }
-      end
-
       def arel_column(columns)
         connection.with_local_table_name { super }
       end
 
       def table_name_matches?(from)
-        connection.with_global_table_name { super }
+        if ::Rails.version < "7.2"
+          connection.with_global_table_name { super }
+        else
+          connection.with_global_table_name do
+            table_name = Regexp.escape(table.name)
+            # INST: adapter_class -> connection
+            quoted_table_name = Regexp.escape(connection.quote_table_name(table.name))
+            /(?:\A|(?<!FROM)\s)(?:\b#{table_name}\b|#{quoted_table_name})(?!\.)/i.match?(from.to_s)
+          end
+        end
+      end
+
+      unless ::Rails.version < "7.2"
+        def order_column(field)
+          arel_column(field) do |attr_name|
+            if attr_name == "count" && !group_values.empty?
+              table[attr_name]
+            else
+              # INST: adapter_class -> connection
+              ::Arel.sql(connection.quote_table_name(attr_name), retryable: true)
+            end
+          end
+        end
       end
 
-      def each_predicate(predicates = nil, &block)
-        return predicates.map(&block) if predicates
+      def each_predicate(predicates = nil, &)
+        return predicates.map(&) if predicates
 
-        each_predicate_cb(:having_clause, :having_clause=, &block)
-        each_predicate_cb(:where_clause, :where_clause=, &block)
+        each_predicate_cb(:having_clause, :having_clause=, &)
+        each_predicate_cb(:where_clause, :where_clause=, &)
       end
 
-      def each_predicate_cb(clause_getter, clause_setter, &block)
+      def each_predicate_cb(clause_getter, clause_setter, &)
         old_clause = send(clause_getter)
         old_predicates = old_clause.send(:predicates)
         return if old_predicates.empty?
 
-        new_predicates = old_predicates.map(&block)
+        new_predicates = old_predicates.map(&)
         return if new_predicates == old_predicates
 
         new_clause = old_clause.dup
@@ -289,7 +327,10 @@ 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)
+            next predicate.class.new predicate.expr.class.new(new_left, new_right) if ::Rails.version < "7.2"
+
+            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
@@ -344,33 +385,33 @@ module Switchman
         end
       end
 
-      def each_transposable_predicate_value_cb(node, original_block, &block)
+      def each_transposable_predicate_value_cb(node, original_block, &)
         case node
         when Array
-          node.filter_map { |val| each_transposable_predicate_value_cb(val, original_block, &block).presence }
+          node.filter_map { |val| each_transposable_predicate_value_cb(val, original_block, &).presence }
         when ::ActiveModel::Attribute
           old_value = node.value_before_type_cast
-          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (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, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (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, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (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, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (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, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           # switch to a regular In, so that Relation::WhereClause#contradiction? knows about it
           if new_value.empty?
@@ -381,7 +422,7 @@ module Switchman
           end
         when ::Arel::Nodes::Binary
           old_value = node.right
-          new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
           (old_value == new_value) ? node : node.class.new(node.left, new_value)
         when ::Arel::Nodes::SelectStatement

data/lib/switchman/active_record/reflection.rb

@@ -28,11 +28,18 @@ module Switchman
         # this technically belongs on AssociationReflection, but we put it on
         # ThroughReflection as well, instead of delegating to its internal
         # HasManyAssociation, losing its proper `klass`
-        def association_scope_cache(klass, owner, &block)
+        def association_scope_cache(klass, owner, &)
           key = self
           key = [key, owner._read_attribute(@foreign_type)] if polymorphic?
           key = [key, shard(owner).id].flatten
-          klass.cached_find_by_statement(key, &block)
+
+          if ::Rails.version < "7.2"
+            klass.cached_find_by_statement(key, &)
+          else
+            klass.with_connection do |connection|
+              klass.cached_find_by_statement(connection, key, &)
+            end
+          end
         end
       end