<commit_before>module ActiveReplicas
  class ProxyingConnection
    def initialize(connection:, is_primary:, proxy:)
      @connection = connection
      @is_primary = is_primary
      @proxy      = proxy
    end

    # role   - Either `:primary` or `:replica`: indicates which database role
    #          is necessary to execute this command.
    # method - Symbol of method to send to a connection.
    def delegate_to(role, method, *args, &block)
      if @is_primary
        @connection.send method, *args, &block
      else
        if role == :primary
          # Need to get a primary connection from the proxy pool.
          @proxy.primary_connection.send method, *args, &block
        else
          @connection.send method, *args, &block
        end
      end
    end

    class << self
      # Partially cribbed from:
      #    https://github.com/kickstarter/replica_pools/blob/master/lib/replica_pools/connection_proxy.rb#L20
      def generate_replica_delegations
        Railtie.replica_delegated_methods.each do |method|
          generate_delegation method, :replica
        end
      end

      def generate_primary_delegations
        Railtie.primary_delegated_methods.each do |method|
          generate_delegation method, :primary
        end
      end

      def generate_delegation(method_name, role)
        class_eval <<-END, __FILE__, __LINE__ + 1
          def #{method_name}(*args, &block)
            delegate_to(:#{role}, :#{method_name}, *args, &block)
          end
        END
      end
    end<commit_msg>Add `with_primary` helper on connection proxy<commit_after>module ActiveReplicas
  class ProxyingConnection
    def initialize(connection:, is_primary:, proxy:)
      @connection = connection
      @is_primary = is_primary
      @proxy      = proxy
    end

    # role   - Either `:primary` or `:replica`: indicates which database role
    #          is necessary to execute this command.
    # method - Symbol of method to send to a connection.
    def delegate_to(role, method, *args, &block)
      if @is_primary
        @connection.send method, *args, &block
      else
        if role == :primary
          # Need to get a primary connection from the proxy pool.
          @proxy.primary_connection.send method, *args, &block
        else
          @connection.send method, *args, &block
        end
      end
    end

    def with_primary
      @proxy.with_primary { yield }
    end

    class << self
      # Partially cribbed from:
      #    https://github.com/kickstarter/replica_pools/blob/master/lib/replica_pools/connection_proxy.rb#L20
      def generate_replica_delegations
        Railtie.replica_delegated_methods.each do |method|
          generate_delegation method, :replica
        end
      end

      def generate_primary_delegations
        Railtie.primary_delegated_methods.each do |method|
          generate_delegation method, :primary
        end
      end

      def generate_delegation(method_name, role)
        class_eval <<-END, __FILE__, __LINE__ + 1
          def #{method_name}(*args, &block)
            delegate_to(:#{role}, :#{method_name}, *args, &block)
          end
        END
      end
    end
{'fraction_non_alphanumeric': 0.07992448080553807, 'fraction_numerical': 0.0018879798615481435, 'mean_word_length': 2.308012486992716, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7138787', 'n_tokens_mistral': 913, 'n_tokens_neox': 917, 'n_words': 251}