<commit_before>require 'thread'
require 'stagehand/active_record_extensions'

module Stagehand
  module Database
    extend self

    def each(&block)
      with_production_connection(&block) unless Configuration.single_connection?
      with_staging_connection(&block)
    end

    def connected_to_production?
      current_connection_name == Configuration.production_connection_name
    end

    def connected_to_staging?
      current_connection_name == Configuration.staging_connection_name
    end

    def production_connection
      ProductionProbe.connection
    end

    def staging_connection
      StagingProbe.connection
    end

    def production_database_name
      database_name(Configuration.production_connection_name)
    end

    def staging_database_name
      database_name(Configuration.staging_connection_name)
    end

    def staging_database_versions
      Stagehand::Database.staging_connection.select_values(versions_scope)
    end

    def production_database_versions
      Stagehand::Database.production_connection.select_values(versions_scope)
    end

    def with_staging_connection(&block)
      with_connection(Configuration.staging_connection_name, &block)
    end

    def with_production_connection(&block)
      with_connection(Configuration.production_connection_name, &block)
    end

    def with_connection(connection_name, &block)
      if current_connection_name != connection_name.to_sym
        Rails.logger.debug "Connecting to #{connection_name}"
        output = swap_connection(connection_name, &block)
        Rails.logger.debug "Restoring connection to #{current_connection_name}"
      else
        Rails.logger.debug "Already connected to #{connection_name}"
        output = yield connection_name
      end
      return output
    end

    def transaction
      success = false
      output = nil
      ActiveRecord::Base.transaction do
        Production::Record.transaction do
          output = yield
          success = true
        end
        raise ActiveRecord::Rollback unless success
      end
      return output
    end

    private

    def swap_connection(connection_name)
      cache = ActiveRecord::Base.connection_pool.query_cache_enabled
      ConnectionStack.push(connection_name.to_sym)
      ActiveRecord::Base.connection_specification_name = current_connection_name
      ActiveRecord::Base.connection_pool.enable_query_cache! if cache

      yield connection_name
    ensure
      ConnectionStack.pop
      ActiveRecord::Base.connection_specification_name = current_connection_name
      ActiveRecord::Base.connection_pool.enable_query_cache! if cache
    end

    def current_connection_name<commit_msg>Don’t pop from connection stack unless pushed<commit_after>require 'thread'
require 'stagehand/active_record_extensions'

module Stagehand
  module Database
    extend self

    def each(&block)
      with_production_connection(&block) unless Configuration.single_connection?
      with_staging_connection(&block)
    end

    def connected_to_production?
      current_connection_name == Configuration.production_connection_name
    end

    def connected_to_staging?
      current_connection_name == Configuration.staging_connection_name
    end

    def production_connection
      ProductionProbe.connection
    end

    def staging_connection
      StagingProbe.connection
    end

    def production_database_name
      database_name(Configuration.production_connection_name)
    end

    def staging_database_name
      database_name(Configuration.staging_connection_name)
    end

    def staging_database_versions
      Stagehand::Database.staging_connection.select_values(versions_scope)
    end

    def production_database_versions
      Stagehand::Database.production_connection.select_values(versions_scope)
    end

    def with_staging_connection(&block)
      with_connection(Configuration.staging_connection_name, &block)
    end

    def with_production_connection(&block)
      with_connection(Configuration.production_connection_name, &block)
    end

    def with_connection(connection_name, &block)
      if current_connection_name != connection_name.to_sym
        Rails.logger.debug "Connecting to #{connection_name}"
        output = swap_connection(connection_name, &block)
        Rails.logger.debug "Restoring connection to #{current_connection_name}"
      else
        Rails.logger.debug "Already connected to #{connection_name}"
        output = yield connection_name
      end
      return output
    end

    def transaction
      success = false
      output = nil
      ActiveRecord::Base.transaction do
        Production::Record.transaction do
          output = yield
          success = true
        end
        raise ActiveRecord::Rollback unless success
      end
      return output
    end

    private

    def swap_connection(connection_name)
      pushed = ConnectionStack.push(connection_name.to_sym)
      cache = ActiveRecord::Base.connection_pool.query_cache_enabled
      ActiveRecord::Base.connection_specification_name = current_connection_name
      ActiveRecord::Base.connection_pool.enable_query_cache! if cache

      yield connection_name
    ensure
      ConnectionStack.pop if pushed
      ActiveRecord::Base.connection_specification_name = current_connection_name
      ActiveRecord::Base.connection_pool.enable_query_cache! if cache
    end

    def current_connection_name
{'fraction_non_alphanumeric': 0.05165098690278547, 'fraction_numerical': 0.0, 'mean_word_length': 3.56013456686291, 'pattern_counts': {'":': 0, '<': 3, '<?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': '4300646', 'n_tokens_mistral': 1401, 'n_tokens_neox': 1359, 'n_words': 284}