<commit_before>module Octoshark
  class ConnectionManager

    attr_reader :connection_pools

    def initialize(configs = {})
      @configs = configs.with_indifferent_access
      setup_connection_pools

      Octoshark.connection_managers << self
    end

    def reset!
      disconnect!
      setup_connection_pools
    end

    def current_connection
      Thread.current[identifier] || raise(Octoshark::Error::NoCurrentConnection, "No current connection")
    end

    def current_connection?
      !Thread.current[identifier].nil?
    end

    def current_or_default_connection
      Thread.current[identifier] || ActiveRecord::Base.connection_pool.connection
    end

    def with_connection(name, &block)
      connection_pool = find_connection_pool(name)
      with_connection_pool(name, connection_pool, &block)
    end

    def with_new_connection(name, config, reusable: false, &block)
      if reusable
        connection_pool = @connection_pools[name] ||= create_connection_pool(config)
        with_connection_pool(name, connection_pool, &block)
      else
        connection_pool = create_connection_pool(config)
        with_connection_pool(name, connection_pool, &block).tap do
          connection_pool.disconnect!
        end
      end
    end

    def without_connection(&block)
      change_connection_reference(nil) do
        yield
      end
    end

    def find_connection_pool(name)
      @connection_pools[name] || raise(Octoshark::Error::NoConnection, "No such database connection '#{name}'")
    end

    def disconnect!
      @connection_pools.values.each do |connection_pool|
        connection_pool.disconnect!
      end
    end

    def identifier
      "octoshark_#{Process.pid}"
    end

    private
    def spec_class
      if defined?(ActiveRecord::ConnectionAdapters::ConnectionSpecification)
        spec_class = ActiveRecord::ConnectionAdapters::ConnectionSpecification
      else
        spec_class = ActiveRecord::Base::ConnectionSpecification
      end
    end

    def change_connection_reference(connection, &block)
      previous_connection = Thread.current[identifier]
      Thread.current[identifier] = connection

      begin
        yield
      ensure
        Thread.current[identifier] = previous_connection
      end
    end

    def setup_connection_pools
      @connection_pools = HashWithIndifferentAccess.new

      @configs.each_pair do |name, config|
        @connection_pools[name] = create_connection_pool(config)
      end<commit_msg>Cache identifier<commit_after>module Octoshark
  class ConnectionManager

    attr_reader :connection_pools

    def initialize(configs = {})
      @configs = configs.with_indifferent_access
      setup_connection_pools

      Octoshark.connection_managers << self
    end

    def reset!
      disconnect!
      setup_connection_pools
    end

    def current_connection
      Thread.current[identifier] || raise(Octoshark::Error::NoCurrentConnection, "No current connection")
    end

    def current_connection?
      !Thread.current[identifier].nil?
    end

    def current_or_default_connection
      Thread.current[identifier] || ActiveRecord::Base.connection_pool.connection
    end

    def with_connection(name, &block)
      connection_pool = find_connection_pool(name)
      with_connection_pool(name, connection_pool, &block)
    end

    def with_new_connection(name, config, reusable: false, &block)
      if reusable
        connection_pool = @connection_pools[name] ||= create_connection_pool(config)
        with_connection_pool(name, connection_pool, &block)
      else
        connection_pool = create_connection_pool(config)
        with_connection_pool(name, connection_pool, &block).tap do
          connection_pool.disconnect!
        end
      end
    end

    def without_connection(&block)
      change_connection_reference(nil) do
        yield
      end
    end

    def find_connection_pool(name)
      @connection_pools[name] || raise(Octoshark::Error::NoConnection, "No such database connection '#{name}'")
    end

    def disconnect!
      @connection_pools.values.each do |connection_pool|
        connection_pool.disconnect!
      end
    end

    def identifier
      @identifier ||= "octoshark_#{Process.pid}"
    end

    private
    def spec_class
      if defined?(ActiveRecord::ConnectionAdapters::ConnectionSpecification)
        spec_class = ActiveRecord::ConnectionAdapters::ConnectionSpecification
      else
        spec_class = ActiveRecord::Base::ConnectionSpecification
      end
    end

    def change_connection_reference(connection, &block)
      previous_connection = Thread.current[identifier]
      Thread.current[identifier] = connection

      begin
        yield
      ensure
        Thread.current[identifier] = previous_connection
      end
    end

    def setup_connection_pools
      @connection_pools = HashWithIndifferentAccess.new

      @configs.each_pair do |name, config|
        @connection_pools[name] = create_connection_pool(config)
      end
{'fraction_non_alphanumeric': 0.07178464606181456, 'fraction_numerical': 0.0, 'mean_word_length': 3.208053691275168, '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': '6420443', 'n_tokens_mistral': 1392, 'n_tokens_neox': 1349, 'n_words': 273}