arget

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function

import sys, os.path, logging, re
from optparse import OptionParser
_log = logging.getLogger('arget')

from twisted.internet import reactor, defer

pname=os.path.basename(sys.argv[0])

act='get'
if pname.startswith('arinfo'):
    act='info'
elif pname.startswith('argrep'):
    act='grep'
elif pname.startswith('arget'):
    act='get'
elif pname.startswith('arsnap'):
    act='snap'
elif pname.startswith('arh5export'):
    act='h5export'

par=OptionParser(
    usage='%prog [options] channel <channels ...>',
    description='Query the channel archiver'
    )

# Select operation
par.add_option('--helptime', action="store_true", default=False,
               help="Show help on date/time format")
par.add_option('--info', action="store_const", dest='act', const='info', default=act,
               help='Show archive server information')
par.add_option('--search', action="store_const", dest='act', const='grep',
               help='Search for channels matching the given pattern(s)')
par.add_option('--get', action="store_const", dest='act', const='get',
               help='Retrieve data for given channels')
par.add_option('--snapshot', action="store_const", dest='act', const='snap',
               help='Retrieve data for given channels')
par.add_option('-E','--export', metavar='TYPE', default=None,
               help="Retrieve data and write to file in the given format (eg. hdf5, pbraw)")

par.add_option('--exact', action='store_const', dest='match', const='exact')
par.add_option('-W','--wildcard', action='store_const', dest='match',
               const='wild')
par.add_option('-R','--regexp', action='store_const', dest='match',
               const='regexp')

# Query options
par.add_option('-s','--start', metavar='TIME',
               help='Start of query window (required)')
par.add_option('-e','--end', metavar='TIME', default=None,
               help='End of query window (defaults to current system time)')
par.add_option('-c','--count', metavar='NUM', default=None, type="int",
               help='Total maximum number of samples to read. (default is inf.)')
par.add_option('-l','--chunk', metavar='NUM', default=1000, type="int",
               help='Maximum number of samples to request in a single query.  (1000 = default)')
par.add_option('-a','--archive', metavar='NAME', action='append', default=[],
               help='Archive name.  Wildcards allowed.  Can be given more than once')
par.add_option('-H','--how', metavar='NAME', default='raw',
               help="Query method (eg. raw or plot)")

# Data options
par.add_option('', '--no-enum', default=False, action='store_true', dest='enumAsInt',
               help='Return enumerations as integers')
par.add_option('', '--skip-first', default=False, action='store_true', dest='skipFirst',
               help="(arget only) Don't print first sample, or PVs with only one sample")

# Input options
par.add_option('', '--pv-list', metavar='FILE', dest='pvlist',
               help='Read PVs from file in addition to argument list ("-" for stdin)')

# Output options
par.add_option('-T','--time', metavar='FMT', dest='timefmt',
               help='Output time format: string, posix')

# General options
par.add_option('-V','--version', action='store_true', help='Show version number')
par.add_option('-v','--verbose', action='count', default=0,
               help='Print more')
par.add_option('-d','--debug', action='store_true', default=False,
               help='Show archiver queries')
par.add_option('-C','--conf', metavar='KEY', default='DEFAULT',
               help='Key for the server configuration')
par.add_option('','--hotshot', metavar='FILE',
               help='Run with the hot-shot profiler.  Write results to the named file')

# PB archive options
par.add_option('--export-no-default-delimiters', action='store_true', help='(pbraw export only) Don\'t use default export delimiters.')
par.add_option('--export-delimiter', metavar='DELIMITER', action='append', help='(pbraw export only) Add extra PV name delimiter.')
par.add_option('--export-granularity', metavar='GRANULARITY', help='(pbraw export only) Time granularity for splitting data into files (5min, 15min, 30min, 1day, 1month, 1year).')
par.add_option('--export-out-dir', metavar='OUT_DIR', help='(pbraw export only) Output directory.')
par.add_option('--appliance-name', metavar='APPLIANCE_NAME', help='(pbraw export only) The name of the appliance to use in mysql.')
par.add_option('--mysql-write-connected', action='store_true', help='(pbraw export only) If defined, mysql statement for the connected and disconnected pvs will be generated, otherwise only disconnected will be included.')

opt, args = par.parse_args()

if opt.helptime:
    from carchive import date
    print(date.__doc__)
    sys.exit(0)

elif opt.version:
    from carchive import __version__
    print(__version__)
    sys.exit(0)

LVL={0:logging.WARN, 1:logging.INFO, 2:logging.DEBUG}

logging.basicConfig(format='%(message)s',level=LVL.get(opt.verbose, LVL[2]))

if opt.pvlist:
    if opt.pvlist == '-':
        AF = sys.stdin
    else:
        AF = open(opt.pvlist, 'r')
    args += filter(lambda L:len(L) and L[0]!='#', map(str.rstrip, AF.readlines()))
    if opt.pvlist != '-':
        AF.close()

from carchive.util import HandledError
from carchive.archive import getArchive
from carchive._conf import loadConfig

conf = loadConfig(opt.conf)

if len(opt.archive)==0:
    opt.archive = conf.get('defaultarchs','*').split()

act = opt.act
if opt.export=='hdf5':
    act='h5export'
elif opt.export=='pbraw':
    act='pbrawexport'

if act=='grep' and not opt.match:
    opt.match='regexp'

@defer.inlineCallbacks
def haveArchive(act, opt, args, conf):

    if opt.verbose>0:
        print('Command:',act)

    mod = __import__('carchive.cmd', fromlist=[act])
    mod = getattr(mod, act)

    if hasattr(mod, 'mangleArgs'):
        opt, args = mod.mangleArgs(opt, args)

    try:
        serv = yield getArchive(conf)
    except HandledError:
        defer.returnValue(None)
    except:
        E = sys.exc_info()[1]
        if opt.verbose<2:
            print('Failed to fetch data server information.',E)
            defer.returnValue(None)
        else:
            raise

    if opt.match=='wild':
        from carchive.util import wild2re
        args = map(wild2re, args)
        opt.match='regexp'

    if opt.verbose>0:
        print('Working with archives', opt.archive)
    opt.archive=serv.archives(pattern=opt.archive)

    breakDown = None
    if opt.match=='regexp':
        breakDown={}
        for arg in args:
            result = yield serv.search(pattern=arg, archs=opt.archive,
                                       breakDown=True, rawTime=True)
            breakDown.update(result)
        args = breakDown.keys()
        if len(args)==0:
            print('Provided pattern(s) did not match any PV')
            defer.returnValue(None)

    elif opt.match!='exact' and len(args)>0:
        breakDown={}
        for arg in args:
            result = yield serv.search(exact=arg, archs=opt.archive,
                                       breakDown=True, rawTime=True)
            breakDown.update(result)
        args = breakDown.keys()
        if len(args)==0:
            print('Provided pattern(s) did not match any PV')
            defer.returnValue(None)

    try:
        import time
        S = time.time()
        yield mod.cmd(action=act, archive=serv,
                      opt=opt, args=args,
                      conf=conf, breakDown=breakDown)
        E = time.time()
        _log.debug('Command runtime %.02f s', E-S)
    except HandledError:
        pass
    except:
        E = sys.exc_info()[1]
        if opt.verbose<2:
            print('Operation failed.',E)
        else:
            raise

@defer.inlineCallbacks
def main(*args):
    try:
        yield haveArchive(*args)
    except:
        logging.exception("Error")
    else:
        reactor.stop()

reactor.callWhenRunning(main, act, opt, args, conf)

if opt.hotshot:
    from hotshot import Profile, stats
    _prof=Profile(opt.hotshot)
    _prof.runcall(reactor.run)
    _prof.close()

    S = stats.load(opt.hotshot)
    S.stream = sys.stderr
    S.sort_stats('time','calls')
    S.print_stats(20)
else:
    reactor.run()

sys.exit(0)