<commit_before>#!/usr/bin/env python

# Copyright 2015 University of Chicago

import sys
import argparse
import logging
import pytz
import os
import pwd
import grp
import datetime

import elasticsearch
import elasticsearch.helpers


ES_NODES = ['uct2-es-head.mwt2.org:9200', 'uct2-es-door.mwt2.org:9200']
TZ_NAME = "US/Central"
TIMEZONE = pytz.timezone(TZ_NAME)


def get_es_client():
    """ Instantiate DB client and pass connection back """
    return elasticsearch.Elasticsearch(hosts=ES_NODES,
                                       retry_on_timeout=True,
                                       max_retries=10,
                                       timeout=300)


def create_record(dirpath, date, index=None):
    """
    Query directory at dirpath for information and store information in ES

    :param dirpath: string storing path to file
    :param date: datetime.date with date to use
    :param inde: ES index for records
    :return: dictionary with record if successful, empty dict otherwise
    """
    if not os.path.isdir(dirpath):
        return {}
    if not index:
        index = "stash-space-{0}".format(date.year, date.month)
    dir_info = os.stat(dirpath)
    uid = dir_info.st_uid
    gid = dir_info.st_gid
    ctime = datetime.datetime.fromtimestamp(dir_info.st_ctime, TIMEZONE)
    mtime = datetime.datetime.fromtimestamp(dir_info.st_mtime, TIMEZONE)
    record_fields = {'@timestamp': date.isoformat(),
                     'size': dir_info.st_size,
                     'ctime': ctime,
                     'mtime': mtime,
                     'user': pwd.getpwuid(uid),
                     'group': grp.getgrgid(gid),
                     'path': dirpath}
    record = {'_index': index,
              '_source': record_fields,
              '_op_type': 'index',
              '_type': 'interval_record'}
    return record


def traverse_directory(dirpath, index=None):
    """
    Traverse subdirectories and create a set of records
    with space usage
    :param dirpath: path to directory to
    :return: Nothing
    """
    current_date = TIMEZONE.localize(datetime.date.today())

    if not os.path.isdir(dirpath):
        return
    records = []
    for root, dirs, files in os.walk(dirpath):
        record = create_record(root, current_date, index)
        if record:
            records.append(record)
    save_records(records)

def save_records(records=None):
    """
    Save a job record to ES
    """
    es_client = get_es_client()
    elasticsearch.helpers.bulk(es_client,
                               records,
                               stats_only=False)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Get space usage on a given '
                                                 'filesystem and store in ES')
    parser.add_argument('--index', dest='index', default=None,
                        help='ES index to store records in')
    parser.add_argument("directory", default=None, required=True,
                        help="Directory to examine")
    args = parser.parse_args(sys.argv[1:])
    if not os.path.isdir(args.directory):
        sys.stderr.write("{0} must be a directory\n".format(args.directory))
        sys.exit(1)
    traverse_directory(args.directory, args.index)<commit_msg>Fix argument addition<commit_after>#!/usr/bin/env python

# Copyright 2015 University of Chicago

import sys
import argparse
import logging
import pytz
import os
import pwd
import grp
import datetime

import elasticsearch
import elasticsearch.helpers


ES_NODES = ['uct2-es-head.mwt2.org:9200', 'uct2-es-door.mwt2.org:9200']
TZ_NAME = "US/Central"
TIMEZONE = pytz.timezone(TZ_NAME)


def get_es_client():
    """ Instantiate DB client and pass connection back """
    return elasticsearch.Elasticsearch(hosts=ES_NODES,
                                       retry_on_timeout=True,
                                       max_retries=10,
                                       timeout=300)


def create_record(dirpath, date, index=None):
    """
    Query directory at dirpath for information and store information in ES

    :param dirpath: string storing path to file
    :param date: datetime.date with date to use
    :param inde: ES index for records
    :return: dictionary with record if successful, empty dict otherwise
    """
    if not os.path.isdir(dirpath):
        return {}
    if not index:
        index = "stash-space-{0}".format(date.year, date.month)
    dir_info = os.stat(dirpath)
    uid = dir_info.st_uid
    gid = dir_info.st_gid
    ctime = datetime.datetime.fromtimestamp(dir_info.st_ctime, TIMEZONE)
    mtime = datetime.datetime.fromtimestamp(dir_info.st_mtime, TIMEZONE)
    record_fields = {'@timestamp': date.isoformat(),
                     'size': dir_info.st_size,
                     'ctime': ctime,
                     'mtime': mtime,
                     'user': pwd.getpwuid(uid),
                     'group': grp.getgrgid(gid),
                     'path': dirpath}
    record = {'_index': index,
              '_source': record_fields,
              '_op_type': 'index',
              '_type': 'interval_record'}
    return record


def traverse_directory(dirpath, index=None):
    """
    Traverse subdirectories and create a set of records
    with space usage
    :param dirpath: path to directory to
    :return: Nothing
    """
    current_date = TIMEZONE.localize(datetime.date.today())

    if not os.path.isdir(dirpath):
        return
    records = []
    for root, dirs, files in os.walk(dirpath):
        record = create_record(root, current_date, index)
        if record:
            records.append(record)
    save_records(records)

def save_records(records=None):
    """
    Save a job record to ES
    """
    es_client = get_es_client()
    elasticsearch.helpers.bulk(es_client,
                               records,
                               stats_only=False)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Get space usage on a given '
                                                 'filesystem and store in ES')
    parser.add_argument('--index', dest='index', default=None,
                        help='ES index to store records in')
    parser.add_argument("directory", default=None, 
                        help="Directory to examine")
    args = parser.parse_args(sys.argv[1:])
    if not os.path.isdir(args.directory):
        sys.stderr.write("{0} must be a directory\n".format(args.directory))
        sys.exit(1)
    traverse_directory(args.directory, args.index)
{'fraction_non_alphanumeric': 0.1003354681305276, 'fraction_numerical': 0.007624275693809088, 'mean_word_length': 2.204201270151441, 'pattern_counts': {'":': 2, '<': 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': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '117763', 'n_tokens_mistral': 1940, 'n_tokens_neox': 1865, 'n_words': 478}