<commit_before>"""
Utility functions used across scripts
"""
__author__ = "<NAME>"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "1.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__date__ = "Dec 1, 2012"

import bson
import datetime
import json
import logging
import os
from pymongo.mongo_client import MongoClient


_log = logging.getLogger("mg.util")

DEFAULT_PORT = 27017
DEFAULT_CONFIG_FILE = "db.json"

DEFAULT_SETTINGS = [
    ("host", "localhost"),
    ("port", DEFAULT_PORT),
    ("database", "vasp"),
    # ("admin_user", None),
    # ("admin_password", None),
    # ("readonly_user", None),
    # ("readonly_password", None),
    # ("collection", "tasks"),
    # ("aliases_config", None),
    # ("mapi_key", None)
]


def get_settings(config_file):
    # Get input file.
    infile = None
    if config_file:
        infile = config_file
    elif os.path.exists(DEFAULT_CONFIG_FILE):
        infile = DEFAULT_CONFIG_FILE

    # Read settings from input file.
    if infile:
        config_settings = json.load(open(infile))
        # Allow some aliases
        for alias, real in (("user", "readonly_user"), ("password", "<PASSWORD>")):
            if alias in config_settings:
                config_settings[real] = config_settings[alias]
                del config_settings[alias]
    else:
        config_settings = {}

    # Merge input file settings and defaults.
    settings = dict(DEFAULT_SETTINGS)
    settings.update(config_settings)

    return settings


def get_database(config_file=None, settings=None, admin=False):
    d = get_settings(config_file) if settings is None else settings
    conn = MongoClient(host=d["host"], port=d["port"])
    db = conn[d["database"]]
    try:
        user = d["admin_user"] if admin else d["readonly_user"]
        passwd = d["admin_password"] if admin else d["readonly_password"]
        db.authenticate(user, passwd)
    except KeyError:
        _log.warn("No {admin,readonly}_user/password found in config. file, "
                  "accessing DB without authentication")
    return db


def get_collection(config_file, admin=False):
    db = get_database(config_file, admin=admin)
    settings = get_settings(config_file)
    return db[settings["collection"]]

<commit_msg>normalize_auth function to help with user/password variations<commit_after>"""
Utility functions used across scripts
"""
__author__ = "<NAME>"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "1.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__date__ = "Dec 1, 2012"

import bson
import datetime
import json
import logging
import os
from pymongo.mongo_client import MongoClient


_log = logging.getLogger("mg.util")

DEFAULT_PORT = 27017
DEFAULT_CONFIG_FILE = "db.json"

DEFAULT_SETTINGS = [
    ("host", "localhost"),
    ("port", DEFAULT_PORT),
    ("database", "vasp"),
    # ("admin_user", None),
    # ("admin_password", None),
    # ("readonly_user", None),
    # ("readonly_password", None),
    # ("collection", "tasks"),
    # ("aliases_config", None),
    # ("mapi_key", None)
]


def get_settings(config_file):
    # Get input file.
    infile = None
    if config_file:
        infile = config_file
    elif os.path.exists(DEFAULT_CONFIG_FILE):
        infile = DEFAULT_CONFIG_FILE

    # Read settings from input file.
    if infile:
        config_settings = json.load(open(infile))
        # Allow some aliases
        for alias, real in (("user", "readonly_user"), ("password", "<PASSWORD>")):
            if alias in config_settings:
                config_settings[real] = config_settings[alias]
                del config_settings[alias]
    else:
        config_settings = {}

    # Merge input file settings and defaults.
    settings = dict(DEFAULT_SETTINGS)
    settings.update(config_settings)

    return settings


def get_database(config_file=None, settings=None, admin=False):
    d = get_settings(config_file) if settings is None else settings
    conn = MongoClient(host=d["host"], port=d["port"])
    db = conn[d["database"]]
    try:
        user = d["admin_user"] if admin else d["readonly_user"]
        passwd = d["admin_password"] if admin else d["readonly_password"]
        db.authenticate(user, passwd)
    except KeyError:
        _log.warn("No {admin,readonly}_user/password found in config. file, "
                  "accessing DB without authentication")
    return db


def normalize_auth(settings, admin=True, readonly=True):
    """Transform the readonly/admin user and password to simple user/password,
    as expected by QueryEngine.

    :param settings: Connection settings
    :type settings: dict
    :param admin: Check for admin password
    :param readonly: Check for readonly password
    :return: Whether user/password are now in settings
    :rtype: bool
    :raise: KeyError if neither readonly, admin, or plain user/password are found
    """
    U, P = "user", "password"
    # If user/password, un-prefixed, exists, do nothing.
    if U in settings and P in settings:
        return True

    # Set prefixes
    prefixes = []
    if admin:
        prefixes.append("admin_")
    if readonly:
        prefixes.append("readonly_")

    # Look for first user/password matching.
    found = False
    for pfx in prefixes:
        ukey, pkey = pfx + U, pfx + P
        if ukey in settings and pkey in settings:
            settings[U] = settings[ukey]
            settings[P] = settings[pkey]
            found = True
            break

    return found


def get_collection(config_file, admin=False):
    db = get_database(config_file, admin=admin)
    settings = get_settings(config_file)
    return db[settings["collection"]]


{'fraction_non_alphanumeric': 0.11328400281888654, 'fraction_numerical': 0.005637773079633545, 'mean_word_length': 3.2020725388601035, 'pattern_counts': {'":': 0, '<': 11, '<?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': '1497956', 'n_tokens_mistral': 1735, 'n_tokens_neox': 1664, 'n_words': 479}