switchman 3.0.1 → 4.2.5

data/lib/switchman/active_record/connection_pool.rb

@@ -1,71 +1,84 @@
 # frozen_string_literal: true
 
-require 'switchman/errors'
-
 module Switchman
   module ActiveRecord
     module ConnectionPool
-      def shard
-        shard_stack.last || Shard.default
-      end
-
-      def shard_stack
-        unless (shard_stack = Thread.current.thread_variable_get(tls_key))
-          shard_stack = Concurrent::Array.new
-          Thread.current.thread_variable_set(tls_key, shard_stack)
-        end
-        shard_stack
-      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
 
       def checkout_new_connection
         conn = super
-        conn.shard = shard
+        conn.shard = current_shard
         conn
       end
 
       def connection(switch_shard: true)
         conn = super()
-        raise NonExistentShardError if shard.new_record?
+        raise Errors::NonExistentShardError if current_shard.new_record?
 
-        switch_database(conn) if conn.shard != shard && switch_shard
+        switch_database(conn) if conn.shard != current_shard && switch_shard
         conn
       end
 
-      def release_connection(with_id = Thread.current)
-        super(with_id)
+      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?
 
-        flush
-      end
+          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?
 
-      def remove_shard!(shard)
-        synchronize do
-          # The shard might be currently active, so we need to update our own shard
-          self.shard = Shard.default if self.shard == shard
-          # Update out any connections that may be using this shard
-          @connections.each do |conn|
-            # This will also update the connection's shard to the default shard
-            switch_database(conn) if conn.shard == shard
+          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
+
+        flush
+      end
+
       def switch_database(conn)
-        @schemas = conn.current_schemas if !@schemas && conn.adapter_name == 'PostgreSQL' && !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 = shard
+        conn.shard = current_shard
       end
 
       private
 
+      def current_shard
+        if ::Rails.version < "8.0"
+          connection_class.current_switchman_shard
+        else
+          connection_descriptor.name.constantize.current_switchman_shard
+        end
+      end
+
       def tls_key
-        "#{object_id}_shard".to_sym
+        :"#{object_id}_shard"
       end
     end
   end

data/lib/switchman/active_record/database_configurations.rb

@@ -3,6 +3,20 @@
 module Switchman
   module ActiveRecord
     module DatabaseConfigurations
+      # key difference: For each env name, ensure only one writable config is returned
+      # 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)
+      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
 
       # key difference: assumes a hybrid two-tier structure; each third tier
@@ -13,19 +27,31 @@ module Switchman
         return configs if configs.is_a?(Array)
 
         db_configs = configs.flat_map do |env_name, config|
-          roles = config.keys.select { |k| config[k].is_a?(Hash) }
-          base_config = config.except(*roles)
+          if config.is_a?(Hash)
+            # 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 do |k|
+              config[k].is_a?(Hash) || (config[k].is_a?(Array) && config[k].all?(Hash))
+            end
+            base_config = config.except(*roles)
+          else
+            base_config = config
+            roles = []
+          end
 
           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]).merge(replica: true))
+            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

@@ -4,10 +4,10 @@ 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_classes)
+          current_shard = Shard.current(klass.connection_class_for_self)
           result = activate do |relation, shard|
             current_id = Shard.relative_id_for(id, current_shard, shard)
             # current_id will be nil for non-integral id
@@ -28,14 +28,14 @@ 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_classes)
+        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)
@@ -43,23 +43,29 @@ module Switchman
       end
 
       def exists?(conditions = :none)
-        conditions = conditions.id if ::ActiveRecord::Base === conditions
-        return false unless conditions
+        return false if @none
 
-        relation = apply_join_dependency(eager_loading: false)
-        return false if ::ActiveRecord::NullRelation === relation
+        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 = relation.except(:select, :order).select('1 AS one').limit(1)
+        return false if !conditions || limit_value == 0 # rubocop:disable Style/NumericPredicate
 
-        case conditions
-        when Array, Hash
-          relation = relation.where(conditions)
-        else
-          relation = relation.where(table[primary_key].eq(conditions)) if conditions != :none
+        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 connection.select_value(shard_rel.arel, "#{name} Exists")
+          return true if skip_query_cache_if_necessary do
+            connection.select_rows(shard_rel.arel, "#{name} Exists?").size == 1
+          end
         end
         false
       end

data/lib/switchman/active_record/log_subscriber.rb

@@ -5,29 +5,26 @@ module Switchman
     module LogSubscriber
       # sadly, have to completely replace this
       def sql(event)
-        self.class.runtime += event.duration
-        return unless logger.debug?
-
         payload = event.payload
 
         return if ::ActiveRecord::LogSubscriber::IGNORE_PAYLOAD_NAMES.include?(payload[:name])
 
         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  = color(sql, sql_color(sql), bold: true)
 
         debug "  #{name}  #{sql}#{binds}#{shard}"
       end

data/lib/switchman/active_record/migration.rb

@@ -14,38 +14,66 @@ module Switchman
 
       def connection
         conn = super
-        ::ActiveRecord::Base.connection_pool.switch_database(conn) if conn.shard != ::ActiveRecord::Base.connection_pool.shard
+        if conn.shard != ::ActiveRecord::Base.current_switchman_shard
+          ::ActiveRecord::Base.connection_pool.switch_database(conn)
+        end
         conn
       end
     end
 
     module Migrator
-      # significant change: hash shard name, not database name
+      # 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
-        shard_name_hash = Zlib.crc32(Shard.current.name)
-        ::ActiveRecord::Migrator::MIGRATOR_SALT * shard_name_hash
+        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
+        @internal_metadata[:migrator_advisory_lock_id] = shard_name_hash.to_s
+        shard_name_hash
       end
 
-      # significant change: strip out prefer_secondary from config
-      def with_advisory_lock_connection
-        pool = ::ActiveRecord::ConnectionAdapters::ConnectionHandler.new.establish_connection(
-          ::ActiveRecord::Base.connection_db_config.configuration_hash.except(:prefer_secondary)
-        )
-
-        pool.with_connection { |connection| yield(connection) } # rubocop:disable Style/ExplicitBlockArgument
-      ensure
-        pool&.disconnect!
+      def use_advisory_lock?
+        super && pending_migrations.any?
       end
     end
 
     module MigrationContext
+      def migrate(...)
+        schema_cache_holder = ::ActiveRecord::Base.connection_pool
+        schema_cache_holder = schema_cache_holder.schema_reflection
+        previous_schema_cache = schema_cache_holder.instance_variable_get(:@cache)
+
+        schema_cache_holder.instance_variable_get(:@cache)
+
+        reset_column_information
+        schema_cache_holder.clear!
+
+        begin
+          super
+        ensure
+          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
+
       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/model_schema.rb

@@ -6,7 +6,7 @@ module Switchman
       module ClassMethods
         def quoted_table_name
           @quoted_table_name ||= {}
-          @quoted_table_name[Shard.current(connection_classes).id] ||= connection.quote_table_name(table_name)
+          @quoted_table_name[Shard.current(connection_class_for_self).id] ||= connection.quote_table_name(table_name)
         end
       end
     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,11 +5,46 @@ module Switchman
     module Persistence
       # touch reads the id attribute directly, so it's not relative to the current shard
       def touch(*, **)
-        shard.activate(self.class.connection_classes) { super }
+        writable_shadow_record_warning
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def update_columns(*)
-        shard.activate(self.class.connection_classes) { super }
+        writable_shadow_record_warning
+        shard.activate(self.class.connection_class_for_self) { super }
+      end
+
+      def delete
+        db = shard.database_server
+        db.unguard { super }
+      end
+
+      def destroy
+        writable_shadow_record_warning
+        super
+      end
+
+      def create_or_update(**, &)
+        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

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
 
@@ -32,12 +32,12 @@ module Switchman
       end
 
       # copy/paste; use quote_local_table_name
-      def drop_database(name) #:nodoc:
+      def drop_database(name) # :nodoc:
         execute "DROP DATABASE IF EXISTS #{quote_local_table_name(name)}"
       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,24 +67,43 @@ 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
 
+      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
+        self.class.quote_table_name(name, shard: @use_local_table_name ? nil : shard)
+      end
 
-        name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
-        name.instance_variable_set(:@schema, shard.name) unless name.schema
-        name.quoted
+      def with_global_table_name(&)
+        with_local_table_name(false, &)
       end
 
       def with_local_table_name(enable = true) # rubocop:disable Style/OptionalBooleanParameter
@@ -97,7 +116,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
 
@@ -124,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/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.send(primary_key)
+            value.send(primary_key(value))
           else
             super
           end

data/lib/switchman/active_record/query_cache.rb

@@ -8,27 +8,36 @@ module Switchman
       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
-          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.2"
+            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
+          else
+            hit = true
+            result = @query_cache.compute_if_absent(key) do
+              hit = false
+              yield
             end
-          result.dup
+          end
         end
+
+        if hit
+          ::ActiveSupport::Notifications.instrument(
+            "sql.active_record",
+            cache_notification_info(sql, name, binds)
+          )
+        end
+
+        result.dup
       end
     end
   end