<commit_before>require 'pry'
require 'active_record'
require 'pg_charmer/connection_handler'

module PgCharmer
  mattr_accessor :default_connection
  mattr_accessor :transaction_connection
  mattr_accessor :force_connection

  # Setup defaults
  self.default_connection = :master
  self.transaction_connection = :master
  self.force_connection = :master if Rails.env.development?

  def self.on(connection)
    self.connection_stack << connection
    begin
      yield
    ensure
      self.connection_stack.pop
    end
  end

  def self.connection_stack
    Thread.current[:pg_charmer_connection_stack] ||= []
  end
end

if Rails::VERSION::MAJOR == 4 and Rails::VERSION::MINOR == 0
  PgCharmer::ConnectionSpecificationResolver = ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver
  ActiveRecord::Base.default_connection_handler = PgCharmer::ConnectionHandler.new
elsif Rails::VERSION::MAJOR == 3 and Rails::VERSION::MINOR == 2
  PgCharmer::ConnectionSpecificationResolver = ActiveRecord::Base::ConnectionSpecification::Resolver
  ActiveRecord::Base.connection_handler = PgCharmer::ConnectionHandler.new
else
  abort "PgCharmer doesn't support Rails version #{Rails.version}"
end

ActiveRecord::Base.instance_eval do
  class_attribute :default_connection

  # disable Model.establish_connection
  def establish_connection(*args)
    # This would be the only place Model.connection_handler.establish_connection is called.
    nil
  end

  # a transaction runs on PgCharmer.connection_for_transactions or on
  # the :on option
  def transaction(options = {}, &block)
    if connection = (options.delete(:on) || PgCharmer.transaction_connection)
      PgCharmer.on(connection) { super }
    else
      super
    end
  end
end<commit_msg>Update docs<commit_after>require 'pry'
require 'active_record'
require 'pg_charmer/connection_handler'

module PgCharmer
  mattr_accessor :default_connection
  mattr_accessor :transaction_connection
  mattr_accessor :force_connection

  # Setup defaults
  self.default_connection = :master
  self.transaction_connection = :master
  self.force_connection = :master if Rails.env.development?

  def self.on(connection)
    self.connection_stack << connection
    begin
      yield
    ensure
      self.connection_stack.pop
    end
  end

  def self.connection_stack
    Thread.current[:pg_charmer_connection_stack] ||= []
  end
end

if Rails::VERSION::MAJOR == 4 and Rails::VERSION::MINOR == 0
  PgCharmer::ConnectionSpecificationResolver = ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver
  ActiveRecord::Base.default_connection_handler = PgCharmer::ConnectionHandler.new
elsif Rails::VERSION::MAJOR == 3 and Rails::VERSION::MINOR == 2
  PgCharmer::ConnectionSpecificationResolver = ActiveRecord::Base::ConnectionSpecification::Resolver
  ActiveRecord::Base.connection_handler = PgCharmer::ConnectionHandler.new
else
  abort "PgCharmer doesn't support Rails version #{Rails.version}"
end

ActiveRecord::Base.instance_eval do
  class_attribute :default_connection

  # disable Model.establish_connection
  def establish_connection(*args)
    # This would be the only place Model.connection_handler.establish_connection is called.
    nil
  end

  # a transaction runs on PgCharmer.transaction_connection or on the :on option
  def transaction(options = {}, &block)
    if connection = (options.delete(:on) || PgCharmer.transaction_connection)
      PgCharmer.on(connection) { super }
    else
      super
    end
  end
end
{'fraction_non_alphanumeric': 0.08335720423947293, 'fraction_numerical': 0.0022916069894013177, 'mean_word_length': 5.224598930481283, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, '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': '1409563', 'n_tokens_mistral': 1023, 'n_tokens_neox': 983, 'n_words': 222}