janus-ar 8.0.0 → 8.0.1

data/lib/janus-ar/active_record/connection_adapters/janus_trilogy_adapter.rb

@@ -1,149 +1,32 @@
-# frozen_string_literal: true
-
-require 'active_record/connection_adapters/abstract_adapter'
-require 'active_record/connection_adapters/trilogy_adapter'
-require_relative '../../../janus-ar'
-
-module ActiveRecord
-  module ConnectionHandling
-    def janus_trilogy_connection(config)
-      ActiveRecord::ConnectionAdapters::JanusTrilogyAdapter.new(config)
-    end
-  end
-end
-
-module ActiveRecord
-  class Base
-    def self.janus_trilogy_adapter_class
-      ActiveRecord::ConnectionAdapters::JanusTrilogyAdapter
-    end
-  end
-end
-
-module ActiveRecord
-  module ConnectionAdapters
-    class JanusTrilogyAdapter < ActiveRecord::ConnectionAdapters::TrilogyAdapter
-      FOUND_ROWS = 'FOUND_ROWS'
-
-      attr_reader :config
-
-      class << self
-        def dbconsole(config, options = {})
-          connection_config = Janus::DbConsoleConfig.new(config)
-
-          super(connection_config, options)
-        end
-      end
-
-      def initialize(*args)
-        args[0][:janus]['replica']['database'] = args[0][:database]
-        args[0][:janus]['primary']['database'] = args[0][:database]
-
-        @replica_config = args[0][:janus]['replica'].symbolize_keys
-        args[0] = args[0][:janus]['primary'].symbolize_keys
-
-        super(*args)
-        @connection_parameters ||= args[0]
-        update_config
-      end
-
-      def raw_execute(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false,
-        materialize_transactions: true, batch: false)
-        case where_to_send?(sql)
-        when :all
-          send_to_replica(sql, connection: :all, method: :execute)
-          super
-        when :replica
-          send_to_replica(sql, connection: :replica, method: :execute)
-        else
-          Janus::Context.stick_to_primary if write_query?(sql)
-          Janus::Context.used_connection(:primary)
-          super
-        end
-      end
-
-      def execute(sql)
-        case where_to_send?(sql)
-        when :all
-          send_to_replica(sql, connection: :all, method: :execute)
-          super(sql)
-        when :replica
-          send_to_replica(sql, connection: :replica, method: :execute)
-        else
-          Janus::Context.stick_to_primary if write_query?(sql)
-          Janus::Context.used_connection(:primary)
-          super(sql)
-        end
-      end
-
-      def execute_and_free(sql, name = nil, async: false)
-        case where_to_send?(sql)
-        when :all
-          send_to_replica(sql, connection: :all, method: :execute)
-          super(sql, name, async:)
-        when :replica
-          send_to_replica(sql, connection: :replica, method: :execute)
-        else
-          Janus::Context.stick_to_primary if write_query?(sql)
-          Janus::Context.used_connection(:primary)
-          super(sql, name, async:)
-        end
-      end
-
-      def with_connection(_args = {})
-        self
-      end
-
-      def connect!(...)
-        replica_connection.connect!(...)
-        super
-      end
-
-      def reconnect!(...)
-        replica_connection.reconnect!(...)
-        super
-      end
-
-      def disconnect!(...)
-        replica_connection.disconnect!(...)
-        super
-      end
-
-      def clear_cache!(...)
-        replica_connection.clear_cache!(...)
-        super
-      end
-
-      def replica_connection
-        @replica_connection ||= ActiveRecord::ConnectionAdapters::TrilogyAdapter.new(@replica_config)
-      end
-
-      private
-
-      def where_to_send?(sql)
-        Janus::QueryDirector.new(sql, open_transactions).where_to_send?
-      end
-
-      def send_to_replica(sql, connection: nil, method: :exec_query)
-        Janus::Context.used_connection(connection) if connection
-        if method == :execute
-          replica_connection.execute(sql)
-        elsif method == :raw_execute
-          replica_connection.execute(sql)
-        else
-          replica_connection.exec_query(sql)
-        end
-      end
-
-      def update_config
-        @config[:flags] ||= 0
-
-        if @config[:flags].is_a? Array
-          @config[:flags].push FOUND_ROWS
-        else
-          @config[:flags] |= ::Janus::Client::FOUND_ROWS
-        end
-      end
-    end
-  end
-end
+# frozen_string_literal: true
+
+require 'active_record/connection_adapters/abstract_adapter'
+require 'active_record/connection_adapters/trilogy_adapter'
+require_relative '../../../janus-ar'
+require_relative '../../adapter_extensions'
+
+module ActiveRecord
+  module ConnectionHandling
+    def janus_trilogy_connection(config)
+      ActiveRecord::ConnectionAdapters::JanusTrilogyAdapter.new(config)
+    end
+  end
+
+  class Base
+    def self.janus_trilogy_adapter_class
+      ActiveRecord::ConnectionAdapters::JanusTrilogyAdapter
+    end
+  end
+
+  module ConnectionAdapters
+    class JanusTrilogyAdapter < ActiveRecord::ConnectionAdapters::TrilogyAdapter
+      include Janus::AdapterExtensions
+
+      private
+
+      def replica_adapter_class
+        ActiveRecord::ConnectionAdapters::TrilogyAdapter
+      end
+    end
+  end
+end

data/lib/janus-ar/adapter_extensions.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+module Janus
+  # Behaviour shared by the Janus MySQL2 and Trilogy adapters.
+  #
+  # Each Janus adapter subclasses the matching ActiveRecord adapter and owns the
+  # *primary* connection (reached via `super`). This module routes every
+  # statement to the primary, a lazily created replica connection, or both, and
+  # keeps Janus::Context up to date.
+  #
+  # Including adapters only need to implement #replica_adapter_class.
+  module AdapterExtensions
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+      def dbconsole(config, options = {})
+        super(Janus::DbConsoleConfig.new(config), options)
+      end
+    end
+
+    attr_reader :config
+
+    def initialize(*args)
+      config = args[0]
+      config[:janus]['replica']['database'] = config[:database]
+      config[:janus]['primary']['database'] = config[:database]
+
+      @replica_config = config[:janus]['replica'].symbolize_keys
+      args[0] = config[:janus]['primary'].symbolize_keys
+
+      super
+      @connection_parameters ||= args[0]
+    end
+
+    # The argument lists below intentionally use anonymous splats and a bare
+    # `super`: ActiveRecord's `raw_execute`/`execute` signatures differ between
+    # versions, so we forward whatever we are given unchanged rather than
+    # restating (and pinning ourselves to) the current signature.
+    def raw_execute(sql, *, **)
+      case where_to_send?(sql)
+      when :all
+        send_to_replica(sql, :all)
+        super
+      when :replica
+        send_to_replica(sql, :replica)
+      else
+        mark_primary(sql)
+        super
+      end
+    end
+
+    def execute(sql, *, **)
+      case where_to_send?(sql)
+      when :all
+        send_to_replica(sql, :all)
+        super
+      when :replica
+        send_to_replica(sql, :replica)
+      else
+        mark_primary(sql)
+        super
+      end
+    end
+
+    def connect!(...)
+      replica_connection.connect!(...)
+      super
+    end
+
+    def reconnect!(...)
+      replica_connection.reconnect!(...)
+      super
+    end
+
+    def disconnect!(...)
+      replica_connection.disconnect!(...)
+      super
+    end
+
+    def clear_cache!(...)
+      replica_connection.clear_cache!(...)
+      super
+    end
+
+    def replica_connection
+      @replica_connection ||= replica_adapter_class.new(@replica_config)
+    end
+
+    private
+
+    def mark_primary(sql)
+      Janus::Context.stick_to_primary if write_query?(sql)
+      Janus::Context.used_connection(:primary)
+    end
+
+    def where_to_send?(sql)
+      Janus::QueryDirector.new(sql, open_transactions).where_to_send?
+    end
+
+    def send_to_replica(sql, connection)
+      Janus::Context.used_connection(connection)
+      replica_connection.execute(sql)
+    end
+  end
+end

data/lib/janus-ar/context.rb

@@ -1,80 +1,79 @@
-# frozen_string_literal: true
-
-module Janus
-  class Context
-    THREAD_KEY = :janus_ar_context
-
-    # Stores the staged data with an expiration time based on the current time,
-    # and clears any expired entries. Returns true if any changes were made to
-    # the current store
-    def initialize(primary: false, expiry: nil)
-      @primary = primary
-      @expiry = expiry
-      @last_used_connection = :primary
-    end
-
-    def stick_to_primary
-      @primary = true
-    end
-
-    def potential_write
-      stick_to_primary
-    end
-
-    def release_all
-      @primary = false
-      @expiry = nil
-      @last_used_connection = nil
-    end
-
-    def use_primary?
-      @primary
-    end
-
-    def used_connection(connection)
-      @last_used_connection = connection
-    end
-
-    attr_reader :last_used_connection
-
-    class << self
-      def stick_to_primary
-        current.stick_to_primary
-      end
-
-      def release_all
-        current.release_all
-      end
-
-      def used_connection(connection)
-        current.used_connection(connection)
-      end
-
-      def use_primary?
-        current.use_primary?
-      end
-
-      def last_used_connection
-        current.last_used_connection
-      end
-
-      protected
-
-      def current
-        fetch(THREAD_KEY) { new }
-      end
-
-      def fetch(key)
-        get(key) || set(key, yield)
-      end
-
-      def get(key)
-        Thread.current.thread_variable_get(key)
-      end
-
-      def set(key, value)
-        Thread.current.thread_variable_set(key, value)
-      end
-    end
-  end
-end
+# frozen_string_literal: true
+
+require 'active_support/isolated_execution_state'
+
+module Janus
+  # Per-execution state that records whether the current unit of work has been
+  # pinned to the primary (e.g. after a write, so subsequent reads stay
+  # consistent). State is stored in ActiveSupport::IsolatedExecutionState so it
+  # follows the application's configured isolation level (thread or fiber),
+  # matching ActiveRecord itself.
+  #
+  # Because pooled threads/fibers are reused across requests and jobs, the
+  # context MUST be released between units of work or a thread that performed a
+  # single write would keep routing every later read to the primary. The Rails
+  # integration (see Janus::Railtie) does this automatically; outside Rails,
+  # call Janus::Context.release_all yourself (e.g. in a Sidekiq middleware).
+  class Context
+    STATE_KEY = :janus_ar_context
+
+    def initialize(primary: false)
+      @primary = primary
+      @last_used_connection = :primary
+    end
+
+    def stick_to_primary
+      @primary = true
+    end
+
+    def release_all
+      @primary = false
+      @last_used_connection = nil
+    end
+
+    def use_primary?
+      @primary
+    end
+
+    def used_connection(connection)
+      @last_used_connection = connection
+    end
+
+    attr_reader :last_used_connection
+
+    class << self
+      def stick_to_primary
+        current.stick_to_primary
+      end
+
+      def release_all
+        current.release_all
+      end
+
+      def used_connection(connection)
+        current.used_connection(connection)
+      end
+
+      def use_primary?
+        current.use_primary?
+      end
+
+      def last_used_connection
+        current.last_used_connection
+      end
+
+      # Release the context at the start of every unit of work wrapped by the
+      # given ActiveSupport executor (web requests, ActiveJob and
+      # Sidekiq-on-Rails jobs all run inside it).
+      def install_reset_hook(executor)
+        executor.to_run { Janus::Context.release_all }
+      end
+
+      protected
+
+      def current
+        ActiveSupport::IsolatedExecutionState[STATE_KEY] ||= new
+      end
+    end
+  end
+end

data/lib/janus-ar/query_director.rb

@@ -1,54 +1,81 @@
-# frozen_string_literal: true
-module Janus
-  class QueryDirector
-    ALL = :all
-    REPLICA = :replica
-    PRIMARY = :primary
-
-    SQL_PRIMARY_MATCHERS = [
-      /\A\s*select.+for update\Z/i, /select.+lock in share mode\Z/i,
-      /\A\s*select.+(nextval|currval|lastval|get_lock|release_lock|pg_advisory_lock|pg_advisory_unlock)\(/i,
-      /\A\s*show/i
-    ].freeze
-    SQL_REPLICA_MATCHERS = [/\A\s*(select|with.+\)\s*select)\s/i].freeze
-    SQL_ALL_MATCHERS = [/\A\s*set\s/i].freeze
-    SQL_SKIP_ALL_MATCHERS = [/\A\s*set\s+local\s/i].freeze
-    WRITE_PREFIXES = %w(INSERT UPDATE DELETE LOCK CREATE GRANT DROP ALTER TRUNCATE BEGIN SAVEPOINT FLUSH).freeze
-
-    def initialize(sql, open_transactions)
-      @_sql = sql
-      @_open_transactions = open_transactions
-    end
-
-    def where_to_send?
-      if should_send_to_all?
-        ALL
-      elsif can_go_to_replica?
-        REPLICA
-      else
-        PRIMARY
-      end
-    end
-
-    private
-
-    def should_send_to_all?
-      SQL_ALL_MATCHERS.any? { |matcher| @_sql =~ matcher } && SQL_SKIP_ALL_MATCHERS.none? { |matcher| @_sql =~ matcher }
-    end
-
-    def can_go_to_replica?
-      !should_go_to_primary?
-    end
-
-    def should_go_to_primary?
-      Janus::Context.use_primary? ||
-        write_query? ||
-        @_open_transactions.positive? ||
-        SQL_PRIMARY_MATCHERS.any? { |matcher| @_sql =~ matcher }
-    end
-
-    def write_query?
-      WRITE_PREFIXES.include?(@_sql.upcase.split(' ').first)
-    end
-  end
-end
+# frozen_string_literal: true
+module Janus
+  class QueryDirector
+    ALL = :all
+    REPLICA = :replica
+    PRIMARY = :primary
+
+    SQL_PRIMARY_MATCHERS = [
+      /\A\s*select.+for update\Z/i, /select.+lock in share mode\Z/i,
+      /\A\s*select.+(nextval|currval|lastval|get_lock|release_lock|pg_advisory_lock|pg_advisory_unlock)\(/i,
+      /\A\s*show/i
+    ].freeze
+    SQL_REPLICA_MATCHERS = [/\A\s*(select|with.+\)\s*select)\s/i].freeze
+    SQL_ALL_MATCHERS = [/\A\s*set\s/i].freeze
+    SQL_SKIP_ALL_MATCHERS = [/\A\s*set\s+local\s/i].freeze
+
+    # Leading whitespace and SQL comments are stripped before matching so that an
+    # annotated statement (e.g. `/* app:web */ INSERT ...`) is classified by the
+    # statement itself rather than by the comment.
+    LEADING_NOISE = %r{\A(?:\s+|/\*.*?\*/|--[^\n]*(?:\n|\z)|\#[^\n]*(?:\n|\z))+}m
+
+    def initialize(sql, open_transactions)
+      @_sql = sql
+      @_open_transactions = open_transactions
+    end
+
+    def where_to_send?
+      if should_send_to_all?
+        ALL
+      elsif can_go_to_replica?
+        REPLICA
+      else
+        PRIMARY
+      end
+    end
+
+    private
+
+    def should_send_to_all?
+      match_any?(SQL_ALL_MATCHERS) && !match_any?(SQL_SKIP_ALL_MATCHERS)
+    end
+
+    # A replica may only serve a statement we positively recognise as a read and
+    # that nothing else forces onto the primary. Everything we do not recognise
+    # as a read defaults to the primary, which is the safe direction for a
+    # write/read proxy: a misrouted read only costs a little primary load, while
+    # a misrouted write is an error (or worse) against a read-only replica.
+    #
+    # Because this is only reached for a confirmed read, there is no need to also
+    # test for a write here.
+    def can_go_to_replica?
+      read_query? && !should_go_to_primary?
+    end
+
+    def read_query?
+      match_any?(SQL_REPLICA_MATCHERS)
+    end
+
+    def should_go_to_primary?
+      Janus::Context.use_primary? ||
+        @_open_transactions.positive? ||
+        match_any?(SQL_PRIMARY_MATCHERS)
+    end
+
+    def match_any?(matchers)
+      matchers.any? { |matcher| normalized_sql.match?(matcher) }
+    end
+
+    # Avoid copying the statement when there is no leading comment/whitespace to
+    # strip, which is the common case for ActiveRecord-generated SQL.
+    def normalized_sql
+      @normalized_sql ||= strip_leading_noise
+    end
+
+    def strip_leading_noise
+      return @_sql unless LEADING_NOISE.match?(@_sql)
+
+      @_sql.sub(LEADING_NOISE, '')
+    end
+  end
+end

data/lib/janus-ar/railtie.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require 'rails/railtie'
+
+module Janus
+  class Railtie < ::Rails::Railtie
+    # Clear Janus' per-request primary stickiness at the start of every unit of
+    # work wrapped by the Rails executor (web requests, ActiveJob and
+    # Sidekiq-on-Rails jobs). Without this, a pooled thread that performs a
+    # write keeps routing reads to the primary for the rest of its life.
+    initializer 'janus.clear_context_per_execution' do |app|
+      Janus::Context.install_reset_hook(app.executor)
+    end
+  end
+end

data/lib/janus-ar/version.rb

@@ -1,17 +1,17 @@
-# frozen_string_literal: true
-
-module Janus
-  unless defined?(::Janus::VERSION)
-    module VERSION
-      MAJOR = 8
-      MINOR = 0
-      PATCH = 0
-      PRE = nil
-
-      def self.to_s
-        [MAJOR, MINOR, PATCH, PRE].compact.join('.')
-      end
-    end
-  end
-  ::Janus::VERSION
-end
+# frozen_string_literal: true
+
+module Janus
+  unless defined?(::Janus::VERSION)
+    module VERSION
+      MAJOR = 8
+      MINOR = 0
+      PATCH = 1
+      PRE = nil
+
+      def self.to_s
+        [MAJOR, MINOR, PATCH, PRE].compact.join('.')
+      end
+    end
+  end
+  ::Janus::VERSION
+end

data/lib/janus-ar.rb

@@ -1,22 +1,24 @@
-# frozen_string_literal: true
-
-require 'active_support'
-
-module Janus
-  autoload :Context, 'janus-ar/context'
-  autoload :Client, 'janus-ar/client'
-  autoload :QueryDirector, 'janus-ar/query_director'
-  autoload :VERSION, 'janus-ar/version'
-  autoload :DbConsoleConfig, 'janus-ar/db_console_config'
-
-  module Logging
-    autoload :Subscriber, 'janus-ar/logging/subscriber'
-    autoload :Logger, 'janus-ar/logging/logger'
-  end
-end
-
-ActiveSupport.on_load(:active_record) do
-  ActiveRecord::LogSubscriber.log_subscribers.each do |subscriber|
-    subscriber.extend Janus::Logging::Subscriber
-  end
-end
+# frozen_string_literal: true
+
+require 'active_support'
+
+module Janus
+  autoload :Context, 'janus-ar/context'
+  autoload :Client, 'janus-ar/client'
+  autoload :QueryDirector, 'janus-ar/query_director'
+  autoload :VERSION, 'janus-ar/version'
+  autoload :DbConsoleConfig, 'janus-ar/db_console_config'
+
+  module Logging
+    autoload :Subscriber, 'janus-ar/logging/subscriber'
+    autoload :Logger, 'janus-ar/logging/logger'
+  end
+end
+
+ActiveSupport.on_load(:active_record) do
+  ActiveRecord::LogSubscriber.log_subscribers.each do |subscriber|
+    subscriber.extend Janus::Logging::Subscriber
+  end
+end
+
+require 'janus-ar/railtie' if defined?(Rails::Railtie)