switchman 2.2.2 → 3.0.0

data/lib/switchman/active_record/connection_handler.rb

@@ -1,190 +0,0 @@
-# frozen_string_literal: true
-
-require 'switchman/connection_pool_proxy'
-
-module Switchman
-  module ActiveRecord
-    module ConnectionHandler
-      def self.make_sharing_automagic(config, shard = Shard.current)
-        # only load the shard name from the db if we have to
-        if !config[:shard_name]
-          # we may not be able to connect to this shard yet, cause it might be an empty database server
-          shard = shard.call if shard.is_a?(Proc)
-          shard_name = shard.name rescue nil
-          return unless shard_name
-
-          config[:shard_name] ||= shard_name
-        end
-      end
-
-      def establish_connection(spec)
-        # Just skip establishing a sharded connection if sharding isn't loaded; we'll do it again later
-        # This only can happen when loading ActiveRecord::Base; after everything is loaded Shard will
-        # be defined and this will actually establish a connection
-        return unless defined?(Shard)
-        pool = super
-
-        # this is the first place that the adapter would have been required; but now we
-        # need this addition ASAP since it will be called when loading the default shard below
-        if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
-          require "switchman/active_record/postgresql_adapter"
-          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter)
-        end
-
-        first_time = !Shard.instance_variable_get(:@default)
-        if first_time
-          # Have to cache the default shard before we insert sharding, otherwise the first access
-          # to sharding will recurse onto itself trying to access column information
-          Shard.default
-
-          config = pool.spec.config
-          # automatically change config to allow for sharing connections with simple config
-          ConnectionHandler.make_sharing_automagic(config)
-          ConnectionHandler.make_sharing_automagic(Shard.default.database_server.config)
-
-          if ::Rails.version < '6.0'
-            ::ActiveRecord::Base.configurations[::Rails.env] = config.stringify_keys
-          else
-            # Adopted from the deprecated code that currently lives in rails proper
-            remaining_configs = ::ActiveRecord::Base.configurations.configurations.reject { |db_config| db_config.env_name == ::Rails.env }
-            new_config = ::ActiveRecord::DatabaseConfigurations.new(::Rails.env => config.stringify_keys).configurations
-            new_configs = remaining_configs + new_config
-
-            ::ActiveRecord::Base.configurations = new_configs
-          end
-        else
-          # this is probably wrong now
-          Shard.default.remove_instance_variable(:@name) if Shard.default.instance_variable_defined?(:@name)
-        end
-
-        @shard_connection_pools ||= { [:primary, Shard.default.database_server.shareable? ? ::Rails.env : Shard.default] => pool}
-
-        category = pool.spec.name.to_sym
-        proxy = ConnectionPoolProxy.new(category,
-                                        pool,
-                                        @shard_connection_pools)
-        owner_to_pool[pool.spec.name] = proxy
-
-        if first_time
-          if Shard.default.database_server.config[:prefer_secondary]
-            Shard.default.database_server.guard!
-          end
-
-          if Shard.default.is_a?(DefaultShard) && Shard.default.database_server.config[:secondary]
-            Shard.default.database_server.guard!
-            Shard.default(reload: true)
-          end
-        end
-
-        # reload the default shard if we just got a new connection
-        # to where the Shards table is
-        # DON'T do it if we're not the current connection handler - that means
-        # we're in the middle of switching environments, and we don't want to
-        # establish a connection with incorrect settings
-        if [:primary, :unsharded].include?(category) && self == ::ActiveRecord::Base.connection_handler && !first_time
-          Shard.default(reload: true, with_fallback: true)
-          proxy.disconnect!
-        end
-
-        if first_time
-          # do the change for other database servers, now that we can switch shards
-          if Shard.default.is_a?(Shard)
-            DatabaseServer.all.each do |server|
-              next if server == Shard.default.database_server
-
-              shard = nil
-              shard_proc = -> do
-                shard ||= server.shards.where(:name => nil).first
-                shard ||= Shard.new(:database_server => server)
-                shard
-              end
-              ConnectionHandler.make_sharing_automagic(server.config, shard_proc)
-              ConnectionHandler.make_sharing_automagic(proxy.current_pool.spec.config, shard_proc)
-            end
-          end
-          # we may have established some connections above trying to infer the shard's name.
-          # close them, so that someone that doesn't expect them doesn't try to fork
-          # without closing them
-          self.clear_all_connections!
-        end
-
-        proxy
-      end
-
-      def remove_connection(spec_name)
-        # also remove pools based on the same spec name that are for shard category purposes
-        # can't just use delete_if, because it's a Concurrent::Map, not a Hash
-        owner_to_pool.keys.each do |k|
-          next if k == spec_name
-
-          v = owner_to_pool[k]
-          owner_to_pool.delete(k) if v.is_a?(ConnectionPoolProxy) && v.default_pool.spec.name == spec_name
-        end
-
-        # unwrap the pool from inside a ConnectionPoolProxy
-        pool = owner_to_pool[spec_name]
-        owner_to_pool[spec_name] = pool.default_pool if pool.is_a?(ConnectionPoolProxy)
-
-        # now let Rails do its thing with the data type it expects
-        super
-      end
-
-      def retrieve_connection_pool(spec_name)
-        owner_to_pool.fetch(spec_name) do
-          if ancestor_pool = pool_from_any_process_for(spec_name)
-            # A connection was established in an ancestor process that must have
-            # subsequently forked. We can't reuse the connection, but we can copy
-            # the specification and establish a new connection with it.
-            spec = if ancestor_pool.is_a?(ConnectionPoolProxy)
-              ancestor_pool.default_pool.spec
-            else
-              ancestor_pool.spec
-            end
-            # avoid copying "duplicate" pools that implement shard categories.
-            # they'll have a spec.name of primary, but a spec_name of something else, like unsharded
-            if spec.name == spec_name
-              pool = establish_connection(spec.to_hash)
-              pool.instance_variable_set(:@schema_cache, ancestor_pool.schema_cache) if ancestor_pool.schema_cache
-              next pool
-            end
-          end
-
-          if spec_name != "primary"
-            primary_pool = retrieve_connection_pool("primary")
-            if primary_pool.is_a?(ConnectionPoolProxy)
-              pool = ConnectionPoolProxy.new(spec_name.to_sym, primary_pool.default_pool, @shard_connection_pools)
-              pool.schema_cache.copy_references(primary_pool.schema_cache)
-              owner_to_pool[spec_name] = pool
-            else
-              primary_pool
-            end
-          else
-            owner_to_pool[spec_name] = nil
-          end
-        end
-      end
-
-      def clear_idle_connections!(since_when)
-        connection_pool_list.each{ |pool| pool.clear_idle_connections!(since_when) }
-      end
-
-      def switchman_connection_pool_proxies
-        owner_to_pool.values.uniq.select{|p| p.is_a?(ConnectionPoolProxy)}
-      end
-
-      private
-
-      # semi-private
-      public
-      def uninitialize_ar(model = ::ActiveRecord::Base)
-        # take the proxies out
-        pool = owner_to_pool[model.name]
-        owner_to_pool[model.name] = pool.default_pool if pool
-      end
-
-      def initialize_categories(model = ::ActiveRecord::Base)
-        class_to_pool.clear
-      end
-    end
-  end
-end

data/lib/switchman/active_record/where_clause_factory.rb

@@ -1,36 +0,0 @@
-# frozen_string_literal: true
-
-module Switchman
-  module ActiveRecord
-    module WhereClauseFactory
-      attr_writer :scope
-
-      def build(opts, other = [])
-        case opts
-        when String, Array
-          values = Hash === other.first ? other.first.values : other
-
-          values.grep(ActiveRecord::Relation) do |rel|
-            # serialize subqueries against the same shard as the outer query is currently
-            # targeted to run against
-            if rel.shard_source_value == :implicit && rel.primary_shard != @scope.primary_shard
-              rel.shard!(@scope.primary_shard)
-            end
-          end
-
-          super
-        when Hash, ::Arel::Nodes::Node
-          where_clause = super
-          binds = ::Rails.version >= "5.2" ? nil : where_clause.binds
-          predicates = where_clause.send(:predicates)
-          @scope.send(:infer_shards_from_primary_key, predicates, binds) if @scope.shard_source_value == :implicit && @scope.shard_value.is_a?(Shard)
-          predicates, _new_binds = @scope.transpose_predicates(predicates, nil, @scope.primary_shard, false, binds: binds)
-          where_clause.instance_variable_set(:@predicates, predicates)
-          where_clause
-        else
-          super
-        end
-      end
-    end
-  end
-end

data/lib/switchman/connection_pool_proxy.rb

@@ -1,173 +0,0 @@
-# frozen_string_literal: true
-
-require 'switchman/schema_cache'
-
-module Switchman
-  module ConnectionError
-    def self.===(other)
-      return true if defined?(PG::Error) && PG::Error === other
-      return true if defined?(Mysql2::Error) && Mysql2::Error === other
-      false
-    end
-  end
-
-  class ConnectionPoolProxy
-    delegate :spec, :connected?, :default_schema, :with_connection, :query_cache_enabled, :active_connection?,
-             :to => :current_pool
-
-    attr_reader :category, :schema_cache
-
-    def default_pool
-      @default_pool
-    end
-
-    def initialize(category, default_pool, shard_connection_pools)
-      @category = category
-      @default_pool = default_pool
-      @connection_pools = shard_connection_pools
-      @schema_cache = default_pool.get_schema_cache(nil) if ::Rails.version >= '6'
-      @schema_cache = SchemaCache.new(self) unless @schema_cache.is_a?(SchemaCache)
-      if ::Rails.version >= '6'
-        @default_pool.set_schema_cache(@schema_cache)
-        @connection_pools.each_value do |pool|
-          pool.set_schema_cache(@schema_cache)
-        end
-      end
-    end
-
-    def active_shard
-      Shard.current(@category)
-    end
-
-    def active_guard_rail_environment
-      ::Rails.env.test? ? :primary : active_shard.database_server.guard_rail_environment
-    end
-
-    def current_pool
-      current_active_shard = active_shard
-      pool = self.default_pool if current_active_shard.database_server == Shard.default.database_server && active_guard_rail_environment == :primary && (current_active_shard.default? || current_active_shard.database_server.shareable?)
-      pool = @connection_pools[pool_key] ||= create_pool unless pool
-      pool.shard = current_active_shard
-      pool
-    end
-
-    def connections
-      connection_pools.map(&:connections).inject([], &:+)
-    end
-
-    def connection(switch_shard: true)
-      pool = current_pool
-      begin
-        connection = pool.connection(switch_shard: switch_shard)
-        connection.instance_variable_set(:@schema_cache, @schema_cache) unless ::Rails.version >= '6'
-        connection
-      rescue ConnectionError
-        raise if active_shard.database_server == Shard.default.database_server && active_guard_rail_environment == :primary
-        configs = active_shard.database_server.config(active_guard_rail_environment)
-        raise unless configs.is_a?(Array)
-        configs.each_with_index do |config, idx|
-          pool = create_pool(config.dup)
-          begin
-            connection = pool.connection
-            connection.instance_variable_set(:@schema_cache, @schema_cache) unless ::Rails.version >= '6'
-          rescue ConnectionError
-            raise if idx == configs.length - 1
-            next
-          end
-          @connection_pools[pool_key] = pool
-          break connection
-        end
-      end
-    end
-
-    def get_schema_cache(_connection)
-      @schema_cache
-    end
-
-    def set_schema_cache(cache)
-      @schema_cache.copy_values(cache)
-    end
-
-    %w{release_connection
-       disconnect!
-       flush!
-       clear_reloadable_connections!
-       verify_active_connections!
-       clear_stale_cached_connections!
-       enable_query_cache!
-       disable_query_cache! }.each do |method|
-      class_eval(<<-RUBY, __FILE__, __LINE__ + 1)
-          def #{method}
-            connection_pools.each(&:#{method})
-          end
-      RUBY
-    end
-
-    def discard!
-      # this breaks everything if i try to pass it onto the pools and i'm not sure why
-    end
-
-    def automatic_reconnect=(value)
-      connection_pools.each { |pool| pool.automatic_reconnect = value }
-    end
-
-    def clear_idle_connections!(since_when)
-      connection_pools.each { |pool| pool.clear_idle_connections!(since_when) }
-    end
-
-    def remove_shard!(shard)
-      connection_pools.each { |pool| pool.remove_shard!(shard) }
-    end
-
-    protected
-
-    def connection_pools
-      (@connection_pools.values + [default_pool]).uniq
-    end
-
-    def pool_key
-      [active_guard_rail_environment,
-        active_shard.database_server.shareable? ? active_shard.database_server.pool_key : active_shard]
-    end
-
-    def create_pool(config = nil)
-      shard = active_shard
-      unless config
-        if shard != Shard.default
-          config = shard.database_server.config(active_guard_rail_environment)
-          config = config.first if config.is_a?(Array)
-          config = config.dup
-        else
-          # we read @config instead of calling config so that we get the config
-          # *before* %{shard_name} is applied
-          # also, we can't just read the database server's config, because
-          # different models could be using different configs on the default
-          # shard, and database server wouldn't know about that
-          config = default_pool.spec.instance_variable_get(:@config)
-          if config[active_guard_rail_environment].is_a?(Hash)
-            config = config.merge(config[active_guard_rail_environment])
-          elsif config[active_guard_rail_environment].is_a?(Array)
-            config = config.merge(config[active_guard_rail_environment].first)
-          else
-            config = config.dup
-          end
-        end
-      end
-      spec = ::ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(
-        default_pool.spec.name,
-        config,
-        "#{config[:adapter]}_connection"
-      )
-      # unfortunately the AR code that does this require logic can't really be
-      # called in isolation
-      require "active_record/connection_adapters/#{config[:adapter]}_adapter"
-
-      ::ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec).tap do |pool|
-        pool.shard = shard
-        pool.set_schema_cache(@schema_cache) if ::Rails.version >= '6'
-        pool.enable_query_cache! if !@connection_pools.empty? && @connection_pools.first.last.query_cache_enabled
-      end
-    end
-  end
-end
-