<commit_before>require 'mongo'
require 'bson'

module Bunsen
  class MongoClient
    include Mongo

    def initialize(database, options = {})
      @options = {
        host: "localhost",
        port: 27017
      }.merge(options)

      @client = Mongo::MongoClient.new(@options[:host], @options[:port], slave_ok: true)
      @database = @client[database]
    end

    def collections
      @database.collections.map(&:name).reject do |collection_name|
        collection_name.start_with?("system.")
      end
    end

    def touch(collection, touch_type = :both)
      should_touch_data = [:data, :both].include?(touch_type.to_sym)
      should_touch_index = [:index, :both].include?(touch_type.to_sym)

      @database.command(
        touch: collection,
        data: should_touch_data,
        index: should_touch_index
      )
    end

    def touch_data_in(collection)
      touch(collection, :data)
    end

    def touch_data_since(collection, time)
      result = {
        "ok" => 0,
        "since" => time,
        "count" => 0
      }

      @database[collection].find(
        _id: { :$gte => time_to_bson_objectid(time) }
      ).each { result["count"] += 1 }

      result["ok"] = 1
      result
    end

    def touch_index_in(collection)
      touch(collection, :index)
    end


    private

    def time_to_bson_objectid(time)
      BSON::ObjectId.from_time(time)
    end
  end
end
<commit_msg>Add a few maintenance mode methods to the DB client<commit_after>require 'mongo'
require 'bson'

module Bunsen
  class MongoClient
    include Mongo

    def initialize(database, options = {})
      @options = {
        host: "localhost",
        port: 27017
      }.merge(options)

      @client = Mongo::MongoClient.new(@options[:host], @options[:port], slave_ok: true)
      @database = @client[database]
      @admin_database = @client["admin"]
    end

    def collections
      @database.collections.map(&:name).reject do |collection_name|
        collection_name.start_with?("system.")
      end
    end

    def enable_maintenance_mode
      @admin_database.command(replSetMaintenance: true)
    end

    def disable_maintenance_mode
      @admin_database.command(replSetMaintenance: false)
    end

    def touch(collection, touch_type = :both)
      should_touch_data = [:data, :both].include?(touch_type.to_sym)
      should_touch_index = [:index, :both].include?(touch_type.to_sym)

      @database.command(
        touch: collection,
        data: should_touch_data,
        index: should_touch_index
      )
    end

    def touch_data_in(collection)
      touch(collection, :data)
    end

    def touch_data_since(collection, time)
      result = {
        "ok" => 0,
        "since" => time,
        "count" => 0
      }

      @database[collection].find(
        _id: { :$gte => time_to_bson_objectid(time) }
      ).each { result["count"] += 1 }

      result["ok"] = 1
      result
    end

    def touch_index_in(collection)
      touch(collection, :index)
    end


    private

    def primary?
      @client.primary?
    end

    def time_to_bson_objectid(time)
      BSON::ObjectId.from_time(time)
    end
  end
end

{'fraction_non_alphanumeric': 0.11367954401519949, 'fraction_numerical': 0.005699810006333122, 'mean_word_length': 2.5454545454545454, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, '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': '5078867', 'n_tokens_mistral': 1038, 'n_tokens_neox': 992, 'n_words': 207}