gridftp_to_es.py

#!/usr/bin/env python
"""
Read from gridftp transfer logs and write to elasticsearch
"""

from __future__ import print_function
import os
import glob
import gzip
import decimal
from optparse import OptionParser
import logging
from functools import partial

parser = OptionParser('usage: %prog [options] transfer_files')
parser.add_option('-a','--address',help='elasticsearch address')
parser.add_option('-n','--indexname',default='gridftp',
                  help='index name (default gridftp)')
(options, args) = parser.parse_args()
if not args:
    parser.error('no gridftp transfer log files')

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(name)s : %(message)s')

def date_convert(d):
    return d[:4]+'-'+d[4:6]+'-'+d[6:8]+'T'+d[8:10]+':'+d[10:12]+':'+d[12:]

def read_from_file(filename):
    with (gzip.open(filename) if filename.endswith('.gz') else open(filename)) as f:
        for line in f:
            try:
                data = {p.split('=',1)[0]:p.split('=',1)[1] for p in line.split() if '=' in p}
                if data['NL.EVNT'] == 'PROG':
                    continue
                data['start_date'] = date_convert(data['START'])
                data['end_date'] = date_convert(data['DATE'])
                data['DEST'] = data['DEST'].strip('[]')
                for k in ('NBYTES','BLOCK','BUFFER','STREAMS'):
                    data[k] = int(data[k])
                data['duration'] = float(decimal.Decimal(data['DATE']) - decimal.Decimal(data['START']))
                data['bandwidth_mbps'] = data['NBYTES'] * 8 / 1000000. / data['duration']
                for k in ('START','DATE','PROG','NL.EVNT','VOLUME','CODE','TASKID'):
                    del data[k]
            except Exception:
                print(data)
                continue
            yield data

def es_generator(entries):
    for data in entries:
        data['_index'] = options.indexname
        data['_type'] = 'transfer_log'
        data['_id'] = data['end_date'].replace('#','-').replace('.','-')
        yield data

from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk

prefix = 'http'
address = options.address
if '://' in address:
    prefix,address = address.split('://')

url = '{}://{}'.format(prefix, address)
logging.info('connecting to ES at %s',url)
es = Elasticsearch(hosts=[url],
                   timeout=5000)
es_import = partial(bulk, es, max_retries=20, initial_backoff=2, max_backoff=3600)

for path in args:
    for filename in glob.iglob(path):
        gen = es_generator(read_from_file(filename))
        success, _ = es_import(gen)
        logging.info('finished processing %s', filename)