<commit_before>class ImportTools::PostgresHandler
  attr_reader :current_conn_values

  def initialize
    self.current_conn_values = ActiveRecord::Base.configurations[Rails.env]
  end

  def with_db db_name, &block
    pg_conn_values = current_conn_values.merge('database' => db_name)
    ActiveRecord::Base.establish_connection pg_conn_values
    block.call(ActiveRecord::Base.connection)
  ensure
    ActiveRecord::Base.establish_connection current_conn_values
  end

  def create_database database_name
    with_db('postgres') { |connection| connection.create_database(database_name) }
  end

  def drop_database database_name
    close_connections_to(database_name)
    with_db('postgres') { |connection| connection.drop_database(database_name) }
  end

  def rename_database database_name, new_database_name
    query = "ALTER DATABASE #{database_name} RENAME TO #{new_database_name}"

    close_connections_to(database_name)
    with_db('postgres') { |connection| connection.execute(query) }
  end

  def seed database_name, dump_path
    pg_password = current_conn_values['password']
    command = []

    command << "PGPASSWORD=#{pg_password}" if pg_password.present?
    command << """
      psql -d #{database_name}
           -U #{current_conn_values['username']}
           -h #{current_conn_values['host']}
      < #{dump_path.to_s}
    """.squish

    system(command.join(" "))
  end

  private
  attr_writer :current_conn_values

  def close_connections_to database_name
    query = """
      SELECT pg_terminate_backend(pg_stat_activity.pid)
      FROM pg_stat_activity
      WHERE pg_stat_activity.datname = '#{database_name}'
        AND pid <> pg_backend_pid();
    """.squish
    with_db('postgres') { |connection| connection.execute(query) }
  end
end
<commit_msg>don't go back to original connection once done<commit_after>class ImportTools::PostgresHandler
  attr_reader :current_conn_values

  def initialize
    self.current_conn_values = ActiveRecord::Base.configurations[Rails.env]
  end

  def with_db db_name, &block
    pg_conn_values = current_conn_values.merge('database' => db_name)
    ActiveRecord::Base.establish_connection pg_conn_values
    block.call(ActiveRecord::Base.connection)
  end

  def create_database database_name
    with_db('postgres') { |connection| connection.create_database(database_name) }
  end

  def drop_database database_name
    close_connections_to(database_name)
    with_db('postgres') { |connection| connection.drop_database(database_name) }
  end

  def rename_database database_name, new_database_name
    query = "ALTER DATABASE #{database_name} RENAME TO #{new_database_name}"

    close_connections_to(database_name)
    with_db('postgres') { |connection| connection.execute(query) }
  end

  def seed database_name, dump_path
    pg_password = current_conn_values['password']
    command = []

    command << "PGPASSWORD=#{pg_password}" if pg_password.present?
    command << """
      psql -d #{database_name}
           -U #{current_conn_values['username']}
           -h #{current_conn_values['host']}
      < #{dump_path.to_s}
    """.squish

    system(command.join(" "))
  end

  private
  attr_writer :current_conn_values

  def close_connections_to database_name
    query = """
      SELECT pg_terminate_backend(pg_stat_activity.pid)
      FROM pg_stat_activity
      WHERE pg_stat_activity.datname = '#{database_name}'
        AND pid <> pg_backend_pid();
    """.squish
    with_db('postgres') { |connection| connection.execute(query) }
  end
end
{'fraction_non_alphanumeric': 0.10025488530161428, 'fraction_numerical': 0.0, 'mean_word_length': 4.327300150829562, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 7, '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': '3823093', 'n_tokens_mistral': 1120, 'n_tokens_neox': 1073, 'n_words': 204}