active_record_proxy_adapters 0.6.1 → 0.7.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d73a2213222b1689c5ff581f40e561974446e480355d582fca5daa04eccf9996
-  data.tar.gz: ee667c46d62816df1dafbe583cbe0503448097f30ebf7c6d37565fd6cb335779
+  metadata.gz: eaa34c675630f13046dbc7170773fca75998eecef0b6e4990ce53110bdfdecef
+  data.tar.gz: 92c024345625535b4f3bfcb241b898e3a26c09836a5f3c6dac61749c913aef85
 SHA512:
-  metadata.gz: a0c1db04e45da543ff86f31a37b45a5f3c5a3e31d0c877cd7e0038f46d01207bf90124868dda9a89ce73a0701726e9c539a422ef53a6e3a9cac9f2304e20a4de
-  data.tar.gz: 1e7eeb8e61c753dd6ad8475a61c4625c90240cdf1d2cf544b40c215a34ec498df806e8403922728c579f5dc727ea5901ff9e81875a57a36fcb2f1fbb17ec4b08
+  metadata.gz: 76493c2922eb8f725410bc69223b9bc9dd121cece620e188a73b941ea7b8ef549108ffea8d3c1c38ddee2ed5686b6160284f87665ea431e41f4a32bccc80472a
+  data.tar.gz: 24c3cf701b469e364a5206ace80d3e8a609fddcd1839337069f0b0f7bbbf27f6c35940d7642a5f717b6030bf776afa988912fb8306b0e8b5d4d2e7fe5bde0ff4

data/README.md

@@ -150,6 +150,25 @@ ActiveRecordProxyAdapters.configure do |config|
 end
 ```
 
+### Configuring multiple connections
+General settings are automatically applied to the `primary` database by default.
+
+If your application has multiple databases that need separate proxy settings, you can use the `database` block for individual database settings.
+
+```ruby
+ActiveRecordProxyAdapters.configure do |config|
+  config.database :primary do |primary_config|
+    primary_config.proxy_delay = 2.seconds
+  end
+
+  config.database :secondary do |secondary_config|
+     secondary_config.proxy_delay = 5.seconds
+  end
+end
+```
+
+With those settings, any model that `connects_to database: { writing: :primary }` will have a 2-second delay, and any model that `connects_to database: { writing: :secondary }` will have a 5-second delay.
+
 ## Logging
 
 ```ruby
@@ -157,8 +176,8 @@ end
 require "active_record_proxy_adapters/log_subscriber"
 
 ActiveRecordProxyAdapters.configure do |config|
-  config.log_subscriber_primary_prefix = "My primary tag" # defaults to "#{adapter_name} Primary", i.e "PostgreSQL Primary"
-  config.log_subscriber_replica_prefix = "My replica tag" # defaults to "#{adapter_name} Replica", i.e "PostgreSQL Replica"
+  config.log_subscriber_primary_prefix = "My primary tag" # defaults to ActiveRecord configuration name", i.e "primary"
+  config.log_subscriber_replica_prefix = "My replica tag" # defaults to ActiveRecord configuration name", i.e "primary_replica"
 end
 
 # You may want to remove duplicate logs

data/lib/active_record_proxy_adapters/active_record_context.rb

@@ -40,7 +40,7 @@ module ActiveRecordProxyAdapters
     def hijackable_methods
       hijackable = %i[execute exec_query]
 
-      hijackable << :internal_exec_query if active_record_v8_0_or_greater?
+      hijackable << :internal_exec_query if active_record_v7_1_or_greater?
 
       hijackable
     end

data/lib/active_record_proxy_adapters/cache_configuration.rb

@@ -7,12 +7,8 @@ module ActiveRecordProxyAdapters
     include SynchronizableConfiguration
 
     # Sets the cache store to use for caching SQL statements.
-    #
-    # @param store [ActiveSupport::Cache::Store] The cache store to use for caching SQL statements.
-    #   Defaults to ActiveSupport::Cache::NullStore, which does not cache anything.
-    #   Thread safe.
+    #   Defaults to {ActiveSupport::Cache::NullStore}, which does not cache anything.
     # @return [ActiveSupport::Cache::Store] The cache store to use for caching SQL statements.
-    #   Defaults to ActiveSupport::Cache::NullStore, which does not cache anything.
     #   Thread safe.
     attr_reader :store
 
@@ -51,5 +47,9 @@ module ActiveRecordProxyAdapters
     def bust
       store.delete_matched("#{key_prefix}*")
     end
+
+    private
+
+    attr_reader :lock
   end
 end

data/lib/active_record_proxy_adapters/configuration.rb

@@ -1,68 +1,66 @@
 # frozen_string_literal: true
 
-require "active_support/core_ext/integer/time"
-require "active_record_proxy_adapters/synchronizable_configuration"
 require "active_record_proxy_adapters/cache_configuration"
+require "active_record_proxy_adapters/context"
+require "active_record_proxy_adapters/database_configuration"
+require "active_record_proxy_adapters/synchronizable_configuration"
+require "active_support/core_ext/integer/time"
 
 module ActiveRecordProxyAdapters
   # Provides a global configuration object to configure how the proxy should behave.
   class Configuration
     include SynchronizableConfiguration
 
-    PROXY_DELAY                   = 2.seconds.freeze
-    CHECKOUT_TIMEOUT              = 2.seconds.freeze
-    LOG_SUBSCRIBER_PRIMARY_PREFIX = proc { |event| "#{event.payload[:connection].class::ADAPTER_NAME} Primary" }.freeze
-    LOG_SUBSCRIBER_REPLICA_PREFIX = proc { |event| "#{event.payload[:connection].class::ADAPTER_NAME} Replica" }.freeze
-
-    # @return [ActiveSupport::Duration] How long the proxy should reroute all read requests to the primary database
-    #   since the latest write. Defaults to PROXY_DELAY. Thread safe.
-    attr_reader :proxy_delay
-    # @return [ActiveSupport::Duration] How long the proxy should wait for a connection from the replica pool.
-    #   Defaults to CHECKOUT_TIMEOUT. Thread safe.
-    attr_reader :checkout_timeout
-
-    # @return [Proc] Prefix for the log subscriber when the primary database is used. Thread safe.
-    attr_reader :log_subscriber_primary_prefix
+    DEFAULT_DATABASE_NAME = :primary
 
-    # @return [Proc] Prefix for the log subscriber when the replica database is used. Thread safe.
-    attr_reader :log_subscriber_replica_prefix
+    # @return [Class] The context that is used to store the current request's state.
+    attr_reader :context_store
 
     def initialize
       @lock = Monitor.new
 
-      self.proxy_delay                   = PROXY_DELAY
-      self.checkout_timeout              = CHECKOUT_TIMEOUT
-      self.log_subscriber_primary_prefix = LOG_SUBSCRIBER_PRIMARY_PREFIX
-      self.log_subscriber_replica_prefix = LOG_SUBSCRIBER_REPLICA_PREFIX
-      self.cache_configuration           = CacheConfiguration.new(@lock)
+      self.cache_configuration = CacheConfiguration.new(@lock)
+      self.context_store       = ActiveRecordProxyAdapters::Context
+      @database_configurations = {}
     end
 
     def log_subscriber_primary_prefix=(prefix)
-      prefix_proc = prefix.is_a?(Proc) ? prefix : proc { prefix.to_s }
+      default_database_config.log_subscriber_primary_prefix = prefix
+    end
 
-      synchronize_update(:log_subscriber_primary_prefix, from: @log_subscriber_primary_prefix, to: prefix_proc) do
-        @log_subscriber_primary_prefix = prefix_proc
-      end
+    def log_subscriber_primary_prefix
+      default_database_config.log_subscriber_primary_prefix
     end
 
     def log_subscriber_replica_prefix=(prefix)
-      prefix_proc = prefix.is_a?(Proc) ? prefix : proc { prefix.to_s }
+      default_database_config.log_subscriber_replica_prefix = prefix
+    end
 
-      synchronize_update(:log_subscriber_replica_prefix, from: @log_subscriber_replica_prefix, to: prefix_proc) do
-        @log_subscriber_replica_prefix = prefix_proc
-      end
+    def log_subscriber_replica_prefix
+      default_database_config.log_subscriber_replica_prefix
+    end
+
+    def proxy_delay
+      default_database_config.proxy_delay
     end
 
     def proxy_delay=(proxy_delay)
-      synchronize_update(:proxy_delay, from: @proxy_delay, to: proxy_delay) do
-        @proxy_delay = proxy_delay
-      end
+      default_database_config.proxy_delay = proxy_delay
+    end
+
+    def checkout_timeout
+      default_database_config.checkout_timeout
     end
 
     def checkout_timeout=(checkout_timeout)
-      synchronize_update(:checkout_timeout, from: @checkout_timeout, to: checkout_timeout) do
-        @checkout_timeout = checkout_timeout
-      end
+      default_database_config.checkout_timeout = checkout_timeout
+    end
+
+    def database(database_name)
+      key = database_name.to_s
+      lock.synchronize { @database_configurations[key] ||= DatabaseConfiguration.new }
+
+      block_given? ? yield(database_configurations[key]) : database_configurations[key]
     end
 
     def cache
@@ -71,13 +69,22 @@ module ActiveRecordProxyAdapters
 
     private
 
-    # @return [CacheConfiguration] The cache configuration for the proxy adapters.
-    attr_reader :cache_configuration
+    attr_reader :cache_configuration, :database_configurations, :lock
+
+    def default_database_config
+      database(DEFAULT_DATABASE_NAME)
+    end
 
     def cache_configuration=(cache_configuration)
       synchronize_update(:cache_configuration, from: @cache_configuration, to: cache_configuration) do
         @cache_configuration = cache_configuration
       end
     end
+
+    def context_store=(context_store)
+      synchronize_update(:context_store, from: @context_store, to: context_store) do
+        @context_store = context_store
+      end
+    end
   end
 end

data/lib/active_record_proxy_adapters/context.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require "active_record_proxy_adapters/mixin/configuration"
+
+module ActiveRecordProxyAdapters
+  # Context is a simple class that holds a registry of connection names and their last write timestamps.
+  # It is used to track the last time a write operation was performed on each connection.
+  # This allows the proxy to determine whether to route read requests to the primary or replica database
+  class Context
+    include Mixin::Configuration
+
+    # @param hash [Hash] A hash containing the connection names as keys and the last write timestamps as values.
+    #   Can be empty.
+    def initialize(hash)
+      @timestamp_registry = hash.transform_values(&:to_f)
+    end
+
+    def recent_write_to_primary?(connection_name)
+      now - self[connection_name.to_s] < proxy_delay(connection_name)
+    end
+
+    def update_for(connection_name)
+      self[connection_name.to_s] = now
+    end
+
+    def [](connection_name)
+      timestamp_registry[connection_name.to_s] || 0
+    end
+
+    def []=(connection_name, timestamp)
+      timestamp_registry[connection_name.to_s] = timestamp
+    end
+
+    def to_h
+      timestamp_registry.dup
+    end
+
+    private
+
+    attr_reader :timestamp_registry
+
+    def now
+      Time.now.utc.to_f
+    end
+  end
+end

data/lib/active_record_proxy_adapters/contextualizer.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module ActiveRecordProxyAdapters
+  # A mixin for managing the context of current database connections.
+  module Contextualizer
+    module_function
+
+    # @return [ActiveRecordProxyAdapters::Context]
+    # Retrieves the context for the current thread.
+    def current_context
+      Thread.current.thread_variable_get(:arpa_context)
+    end
+
+    # @param context [ActiveRecordProxyAdapters::Context]
+    # Sets the context for the current thread.
+    def current_context=(context)
+      Thread.current.thread_variable_set(:arpa_context, context)
+    end
+  end
+end

data/lib/active_record_proxy_adapters/database_configuration.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/integer/time"
+require "active_record_proxy_adapters/synchronizable_configuration"
+require "active_record_proxy_adapters/cache_configuration"
+require "active_record_proxy_adapters/context"
+
+module ActiveRecordProxyAdapters
+  # Provides a global configuration object to configure how the proxy should behave.
+  class DatabaseConfiguration
+    include SynchronizableConfiguration
+
+    PROXY_DELAY      = 2.seconds.freeze
+    CHECKOUT_TIMEOUT = 2.seconds.freeze
+    DEFAULT_PREFIX   = proc do |event|
+      connection = event.payload[:connection]
+
+      connection.pool.try(:db_config).try(:name) || connection.class::ADAPTER_NAME
+    end
+
+    # @return [ActiveSupport::Duration] How long the proxy should reroute all read requests to the primary database
+    #   since the latest write. Defaults to PROXY_DELAY. Thread safe.
+    attr_reader :proxy_delay
+    # @return [ActiveSupport::Duration] How long the proxy should wait for a connection from the replica pool.
+    #   Defaults to CHECKOUT_TIMEOUT. Thread safe.
+    attr_reader :checkout_timeout
+
+    # @return [Proc] Prefix for the log subscriber when the primary database is used. Thread safe.
+    attr_reader :log_subscriber_primary_prefix
+
+    # @return [Proc] Prefix for the log subscriber when the replica database is used. Thread safe.
+    attr_reader :log_subscriber_replica_prefix
+
+    def initialize
+      @lock                              = Monitor.new
+      self.proxy_delay                   = PROXY_DELAY
+      self.checkout_timeout              = CHECKOUT_TIMEOUT
+      self.log_subscriber_primary_prefix = DEFAULT_PREFIX
+      self.log_subscriber_replica_prefix = DEFAULT_PREFIX
+    end
+
+    def proxy_delay=(proxy_delay)
+      synchronize_update(:proxy_delay, from: @proxy_delay, to: proxy_delay) do
+        @proxy_delay = proxy_delay
+      end
+    end
+
+    def checkout_timeout=(checkout_timeout)
+      synchronize_update(:checkout_timeout, from: @checkout_timeout, to: checkout_timeout) do
+        @checkout_timeout = checkout_timeout
+      end
+    end
+
+    def log_subscriber_primary_prefix=(prefix)
+      prefix_proc = prefix.is_a?(Proc) ? prefix : proc { prefix.to_s }
+
+      synchronize_update(:log_subscriber_primary_prefix, from: @log_subscriber_primary_prefix, to: prefix_proc) do
+        @log_subscriber_primary_prefix = prefix_proc
+      end
+    end
+
+    def log_subscriber_replica_prefix=(prefix)
+      prefix_proc = prefix.is_a?(Proc) ? prefix : proc { prefix.to_s }
+
+      synchronize_update(:log_subscriber_replica_prefix, from: @log_subscriber_replica_prefix, to: prefix_proc) do
+        @log_subscriber_replica_prefix = prefix_proc
+      end
+    end
+
+    private
+
+    attr_reader :lock
+  end
+end

data/lib/active_record_proxy_adapters/log_subscriber.rb

@@ -2,6 +2,8 @@
 
 module ActiveRecordProxyAdapters
   class LogSubscriber < ActiveRecord::LogSubscriber # rubocop:disable Style/Documentation
+    include Mixin::Configuration
+
     attach_to :active_record
 
     IGNORE_PAYLOAD_NAMES = %w[SCHEMA EXPLAIN].freeze
@@ -19,23 +21,27 @@ module ActiveRecordProxyAdapters
     protected
 
     def database_instance_prefix_for(event)
-      connection = event.payload[:connection]
-      config = connection.instance_variable_get(:@config)
-      prefix = if config[:replica] || config["replica"]
-                 log_subscriber_replica_prefix
+      connection      = event.payload[:connection]
+      db_config       = connection.pool.try(:db_config) || NullConfig.new # AR 7.0.x does not support "NullConfig"
+      connection_name = db_config.name
+
+      prefix = if db_config.replica?
+                 log_subscriber_replica_prefix(connection_name)
                else
-                 log_subscriber_primary_prefix
+                 log_subscriber_primary_prefix(connection_name)
                end
 
       "[#{prefix.call(event)}]"
     end
 
-    private
-
-    delegate :log_subscriber_primary_prefix, :log_subscriber_replica_prefix, to: :config
+    class NullConfig # rubocop:disable Style/Documentation
+      def method_missing(...)
+        nil
+      end
 
-    def config
-      ActiveRecordProxyAdapters.config
+      def respond_to_missing?(*)
+        true
+      end
     end
   end
 end

data/lib/active_record_proxy_adapters/middleware.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require "rack"
+require "json"
+require "active_record_proxy_adapters/context"
+require "active_record_proxy_adapters/contextualizer"
+require "active_record_proxy_adapters/mixin/configuration"
+
+module ActiveRecordProxyAdapters
+  class Middleware # rubocop:disable Style/Documentation
+    include Contextualizer
+    include Mixin::Configuration
+
+    COOKIE_NAME   = "arpa_context"
+    COOKIE_BUFFER = 5.seconds.freeze
+    DEFAULT_COOKIE_OPTIONS = {
+      path: "/",
+      http_only: true
+    }.freeze
+
+    class << self
+      include Mixin::Configuration
+    end
+
+    COOKIE_READER = lambda do |rack_env|
+      rack_request = Rack::Request.new(rack_env)
+      arpa_cookie  = rack_request.cookies[COOKIE_NAME]
+      JSON.parse(arpa_cookie || "{}")
+    rescue JSON::ParserError
+      {}
+    end.freeze
+
+    COOKIE_WRITER = lambda do |headers, cookie_hash, options|
+      cookie           = DEFAULT_COOKIE_OPTIONS.merge(options)
+      max_value        = cookie_hash.values.max || 0
+      then_time        = Time.at(max_value).utc
+      expires          = then_time + proxy_delay(cookie_hash.key(max_value)) + COOKIE_BUFFER
+      max_age          = expires - then_time
+      cookie[:expires] = expires
+      cookie[:max_age] = max_age
+      cookie[:value]   = cookie_hash.to_json
+
+      Rack::Utils.set_cookie_header!(headers, COOKIE_NAME, cookie)
+    end.freeze
+
+    def initialize(app, cookie_options = {})
+      @app = app
+      @cookie_options = cookie_options
+    end
+
+    def call(env)
+      return @app.call(env) if ignore_request?(env)
+
+      self.current_context = context_store.new(COOKIE_READER.call(env))
+
+      status, headers, body = @app.call(env)
+
+      update_cookie_from_context(headers)
+
+      [status, headers, body]
+    end
+
+    private
+
+    def update_cookie_from_context(headers)
+      COOKIE_WRITER.call(headers, current_context.to_h, @cookie_options)
+    end
+
+    def ignore_request?(env)
+      return false unless defined?(Rails)
+      return false unless asset_prefix
+
+      /^#{asset_prefix}/.match?(env["PATH_INFO"].to_s)
+    end
+
+    def asset_prefix
+      Rails.try(:application).try(:config).try(:assets).try(:prefix)
+    end
+  end
+end

data/lib/active_record_proxy_adapters/mixin/configuration.rb

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/integer/time"
+require "active_record_proxy_adapters/synchronizable_configuration"
+require "active_record_proxy_adapters/cache_configuration"
+require "active_record_proxy_adapters/context"
+
+module ActiveRecordProxyAdapters
+  module Mixin
+    # Provides helpers to access to reduce boilerplate while retrieving configuration properties.
+    module Configuration
+      # Helper to retrieve the proxy delay from the configuration stored in
+      # {ActiveRecordProxyAdapters::DatabaseConfiguration#log_subscriber_primary_prefix}.
+      # @param database_name [Symbol, String] The name of the database to retrieve the prefix.
+      # @return [Proc]
+      def log_subscriber_primary_prefix(database_name)
+        database_config(database_name).log_subscriber_primary_prefix
+      end
+
+      # Helper to retrieve the proxy delay from the configuration stored in
+      # {ActiveRecordProxyAdapters::DatabaseConfiguration#log_subscriber_replica_prefix}.
+      # @param database_name [Symbol, String] The name of the database to retrieve the prefix.
+      # @return [Proc]
+      def log_subscriber_replica_prefix(database_name)
+        database_config(database_name).log_subscriber_replica_prefix
+      end
+
+      # Helper to retrieve the proxy delay from the configuration stored in
+      # {ActiveRecordProxyAdapters::DatabaseConfiguration#proxy_delay}.
+      # @param database_name [Symbol, String] The name of the database to retrieve the proxy delay for.
+      # @return [ActiveSupport::Duration]
+      def proxy_delay(database_name)
+        database_config(database_name).proxy_delay
+      end
+
+      # Helper to retrieve the checkout timeout from the configuration stored in
+      # {ActiveRecordProxyAdapters::DatabaseConfiguration#checkout_timeout}.
+      # @param database_name [Symbol, String] The name of the database to retrieve the checkout timeout for.
+      # @return [ActiveSupport::Duration]
+      def checkout_timeout(database_name)
+        database_config(database_name).checkout_timeout
+      end
+
+      # Helper to retrieve the context store class from the configuration stored in
+      # {ActiveRecordProxyAdapters::Configuration#context_store}.
+      # @return [Class]
+      def context_store
+        proxy_config.context_store
+      end
+
+      # Helper to retrieve the cache store from the configuration stored in
+      # {ActiveRecordProxyAdapters::CacheConfiguration#store}.
+      # @return [ActiveSupport::Cache::Store]
+      def cache_store
+        cache_config.store
+      end
+
+      # Helper to retrieve the cache key prefix from the configuration stored in
+      # {ActiveRecordProxyAdapters::CacheConfiguration#key_prefix}.
+      # It uses the key builder to generate a cache key for the given SQL string, prepended with the key prefix.
+      # @return [String]
+      def cache_key_for(sql_string)
+        cache_config.key_builder.call(sql_string).prepend(cache_config.key_prefix)
+      end
+
+      # @!visibility private
+      def cache_config
+        proxy_config.cache
+      end
+
+      # @!visibility private
+      def database_config(database_name)
+        proxy_config.database(database_name)
+      end
+
+      # @!visibility private
+      def proxy_config
+        ActiveRecordProxyAdapters.config
+      end
+    end
+  end
+end

data/lib/active_record_proxy_adapters/primary_replica_proxy.rb

@@ -1,17 +1,20 @@
 # frozen_string_literal: true
 
+require "active_record_proxy_adapters/active_record_context"
 require "active_record_proxy_adapters/configuration"
+require "active_record_proxy_adapters/contextualizer"
+require "active_record_proxy_adapters/hijackable"
+require "active_record_proxy_adapters/mixin/configuration"
 require "active_support/core_ext/module/delegation"
 require "active_support/core_ext/object/blank"
-require "concurrent-ruby"
-require "active_record_proxy_adapters/active_record_context"
-require "active_record_proxy_adapters/hijackable"
 
 module ActiveRecordProxyAdapters
   # This is the base class for all proxies. It defines the methods that should be proxied
   # and the logic to determine which database to use.
   class PrimaryReplicaProxy # rubocop:disable Metrics/ClassLength
     include Hijackable
+    include Contextualizer
+    include Mixin::Configuration
 
     # All queries that match these patterns should be sent to the primary database
     SQL_PRIMARY_MATCHERS = [
@@ -51,13 +54,12 @@ module ActiveRecordProxyAdapters
     # @param primary_connection [ActiveRecord::ConnectionAdatpers::AbstractAdapter]
     def initialize(primary_connection)
       @primary_connection    = primary_connection
-      @last_write_at         = 0
       @active_record_context = ActiveRecordContext.new
     end
 
     private
 
-    attr_reader :primary_connection, :last_write_at, :active_record_context
+    attr_reader :primary_connection, :active_record_context
 
     delegate :connection_handler, to: :connection_class
     delegate :reading_role, :writing_role, to: :active_record_context
@@ -140,10 +142,6 @@ module ActiveRecordProxyAdapters
       [reading_role, writing_role].include?(role) ? role : nil
     end
 
-    def cache_key_for(sql_string)
-      cache_config.key_builder.call(sql_string).prepend(cache_config.key_prefix)
-    end
-
     def connected_to_stack
       return connection_class.connected_to_stack if connection_class.respond_to?(:connected_to_stack)
 
@@ -183,7 +181,7 @@ module ActiveRecordProxyAdapters
     end
 
     def checkout_replica_connection
-      replica_pool.checkout(checkout_timeout)
+      replica_pool.checkout(checkout_timeout(primary_connection_name))
     # rescue NoDatabaseError to avoid crashing when running db:create rake task
     # rescue ConnectionNotEstablished to handle connectivity issues in the replica
     # (for example, replication delay)
@@ -191,7 +189,6 @@ module ActiveRecordProxyAdapters
       primary_connection
     end
 
-    # @return [TrueClass] if there has been a write within the last {#proxy_delay} seconds
     # @return [TrueClass] if sql_string matches a write statement (i.e. INSERT, UPDATE, DELETE, SELECT FOR UPDATE)
     # @return [FalseClass] if sql_string matches a read statement (i.e. SELECT)
     def need_primary?(sql_string)
@@ -221,11 +218,9 @@ module ActiveRecordProxyAdapters
       proc { |matcher| matcher.match?(sql_string) }
     end
 
-    # TODO: implement a context API to re-route requests to the primary database if a recent query was sent to it to
-    # avoid replication delay issues
     # @return Boolean
     def recent_write_to_primary?
-      Concurrent.monotonic_time - last_write_at < proxy_delay
+      proxy_context.recent_write_to_primary?(primary_connection_name)
     end
 
     def in_transaction?
@@ -233,27 +228,15 @@ module ActiveRecordProxyAdapters
     end
 
     def update_primary_latest_write_timestamp
-      @last_write_at = Concurrent.monotonic_time
-    end
-
-    def cache_store
-      cache_config.store
-    end
-
-    def cache_config
-      proxy_config.cache
-    end
-
-    def proxy_delay
-      proxy_config.proxy_delay
+      proxy_context.update_for(primary_connection_name)
     end
 
-    def checkout_timeout
-      proxy_config.checkout_timeout
+    def primary_connection_name
+      @primary_connection_name ||= primary_connection.pool.try(:db_config).try(:name).try(:to_s)
     end
 
-    def proxy_config
-      ActiveRecordProxyAdapters.config
+    def proxy_context
+      self.current_context ||= context_store.new({})
     end
   end
 end

data/lib/active_record_proxy_adapters/railtie.rb

@@ -6,6 +6,7 @@ module ActiveRecordProxyAdapters
   # Hooks into rails boot process to extend ActiveRecord with the proxy adapter.
   class Railtie < Rails::Railtie
     require "active_record_proxy_adapters/connection_handling"
+    require "active_record_proxy_adapters/middleware"
 
     config.to_prepare do
       Rails.autoloaders.each do |autoloader|
@@ -17,5 +18,9 @@ module ActiveRecordProxyAdapters
         )
       end
     end
+
+    initializer "active_record_proxy_adapters.configure_rails_initialization" do |app|
+      app.middleware.use ActiveRecordProxyAdapters::Middleware
+    end
   end
 end

data/lib/active_record_proxy_adapters/synchronizable_configuration.rb

@@ -7,8 +7,6 @@ module ActiveRecordProxyAdapters
     included do
       private
 
-      attr_reader :lock
-
       def synchronize_update(attribute, from:, to:, &block)
         ActiveSupport::Notifications.instrument(
           "active_record_proxy_adapters.configuration_update",

data/lib/active_record_proxy_adapters/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module ActiveRecordProxyAdapters
-  VERSION = "0.6.1"
+  VERSION = "0.7.0"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: active_record_proxy_adapters
 version: !ruby/object:Gem::Version
-  version: 0.6.1
+  version: 0.7.0
 platform: ruby
 authors:
 - Matt Cruz
@@ -63,6 +63,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: 3.1.0
+- !ruby/object:Gem::Dependency
+  name: json
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 description: |-
   This gem allows automatic connection switching between a primary and one read replica database in ActiveRecord.
   It pattern matches the SQL statement being sent to decide whether it should go to the replica (SELECT) or the
@@ -94,9 +108,14 @@ files:
 - lib/active_record_proxy_adapters/connection_handling/postgresql.rb
 - lib/active_record_proxy_adapters/connection_handling/sqlite3.rb
 - lib/active_record_proxy_adapters/connection_handling/trilogy.rb
+- lib/active_record_proxy_adapters/context.rb
+- lib/active_record_proxy_adapters/contextualizer.rb
+- lib/active_record_proxy_adapters/database_configuration.rb
 - lib/active_record_proxy_adapters/database_tasks.rb
 - lib/active_record_proxy_adapters/hijackable.rb
 - lib/active_record_proxy_adapters/log_subscriber.rb
+- lib/active_record_proxy_adapters/middleware.rb
+- lib/active_record_proxy_adapters/mixin/configuration.rb
 - lib/active_record_proxy_adapters/mysql2_proxy.rb
 - lib/active_record_proxy_adapters/postgresql_proxy.rb
 - lib/active_record_proxy_adapters/primary_replica_proxy.rb