nilmdb/nilmdb/cmdline.py

"""Command line client functionality, broken into a separate file
so it can be more easily tested."""

from __future__ import absolute_import
from nilmdb.printf import *
import nilmdb.client
import nilmdb.layout
import nilmdb.timestamper

import time
import sys
import re
import os
import urlparse
import argparse
import fnmatch
import subprocess

from argparse import ArgumentDefaultsHelpFormatter as def_form

version = "0.1"

class Cmdline(object):

    def __init__(self, argv):
        self.argv = argv

    def parser_setup(self):
        version_string = sprintf("nilmtool %s, client library %s",
                                 version, nilmdb.Client.client_version)

        self.parser = argparse.ArgumentParser(add_help = False,
                                              formatter_class = def_form)

        group = self.parser.add_argument_group("General options")
        group.add_argument("-q", "--quiet", action='store_true',
                           help='suppress unnecessary messages')
        group.add_argument("-h", "--help", action='help',
                           help='show this help message and exit')
        group.add_argument("-V", "--version", action="version",
                           version=version_string)

        group = self.parser.add_argument_group("Server")
        group.add_argument("-u", "--url", action="store",
                           default="http://localhost:12380/",
                           help="NilmDB server URL (default: %(default)s)")

        sub = self.parser.add_subparsers(title="Commands",
                                         dest="command",
                                         description="Specify --help after "
                                         "the command for command-specific "
                                         "options.")

        self.parser_setup_info(sub)
        self.parser_setup_list(sub)
        self.parser_setup_create(sub)
        self.parser_setup_metadata(sub)
        self.parser_setup_insert(sub)

    def parser_setup_info(self, sub):
        cmd = sub.add_parser("info", help="Server information",
                             formatter_class = def_form,
                             description="""
                             List information about the server, like
                             version.
                             """)
        cmd.set_defaults(handler = self.cmd_info)

    def parser_setup_list(self, sub):
        cmd = sub.add_parser("list", help="List streams",
                             formatter_class = def_form,
                             description="""
                             List streams available in the database,
                             optionally filtering by layout or path.  Wildcards
                             are accepted.
                             """)
        cmd.set_defaults(handler = self.cmd_list)
        group = cmd.add_argument_group("Stream filtering")
        group.add_argument("-l", "--layout", default="*",
                           help="Match only this stream layout")
        group.add_argument("-p", "--path", default="*",
                           help="Match only this path")

    def parser_setup_create(self, sub):
        cmd = sub.add_parser("create", help="Create a new stream",
                             formatter_class = def_form,
                             description="""
                             Create a new empty stream at the
                             specified path and with the specifed
                             layout type.
                             """)
        cmd.set_defaults(handler = self.cmd_create)
        group = cmd.add_argument_group("Required arguments")
        group.add_argument("path",
                           help="Path of new stream, e.g. /foo/bar")
        group.add_argument("layout",
                           help="Layout type for new stream, e.g. RawData")

    def parser_setup_metadata(self, sub):
        cmd = sub.add_parser("metadata", help="Get or set stream metadata",
                             description="""
                             Get or set key=value metadata associated with
                             a stream.
                             """,
                             usage="%(prog)s path [-g [key ...] | "
                             "-s key=value [...] | -u key=value [...]]")
        cmd.set_defaults(handler = self.cmd_metadata)

        group = cmd.add_argument_group("Required arguments")
        group.add_argument("path",
                           help="Path of stream, e.g. /foo/bar")

        group = cmd.add_argument_group("Actions")
        exc = group.add_mutually_exclusive_group()
        exc.add_argument("-g", "--get", nargs="*", metavar="key",
                         help="Get metadata for specified keys (default all)")
        exc.add_argument("-s", "--set", nargs="+", metavar="key=value",
                         help="Replace all metadata with provided "
                         "key=value pairs")
        exc.add_argument("-u", "--update", nargs="+", metavar="key=value",
                         help="Update metadata using provided "
                         "key=value pairs")

    def parser_setup_insert(self, sub):
        cmd = sub.add_parser("insert", help="Insert data",
                             description="""
                             Insert data into a stream.
                             """)
        cmd.set_defaults(handler = self.cmd_insert)

        group = cmd.add_argument_group("Timestamping",
                                       description="""
                                       If timestamps are already provided in the
                                       input date, use --none.  Otherwise,
                                       provide --start, or use --filename to
                                       try to deduce timestamps from the file.
                                       """)

        group.add_argument("-r", "--rate", type=float,
                           help="""
                           If needed, rate in Hz (default: based on
                           stream layout)
                           """)
        exc = group.add_mutually_exclusive_group()
        exc.add_argument("-s", "--start", metavar="TIME",
                         help="Starting timestamp (free-form)")
        exc.add_argument("-f", "--filename", action="store_true",
                         help="""
                         Use filenames to determine start time
                         (default, if filenames are provided)
                         """)
        exc.add_argument("-n", "--none", action="store_true",
                         help="Timestamp is already present, don't add one")

        group = cmd.add_argument_group("Required parameters")
        group.add_argument("path",
                           help="Path of stream, e.g. /foo/bar")
        group.add_argument("file", nargs="*", default=['-'],
                           help="File(s) to insert (default: stdin)")

    def die(self, formatstr, *args):
        fprintf(sys.stderr, formatstr, *args)
        self.client.close()
        sys.exit(-1)

    def run(self):
        # Run parser
        self.parser_setup()
        self.args = self.parser.parse_args(self.argv)

        self.client = nilmdb.Client(self.args.url)

        # Make a test connection to make sure things work
        try:
            server_version = self.client.version()
        except nilmdb.client.NilmCommError as e:
            self.die("Error connecting to server: %s\n", str(e))

        # Now dispatch client request to appropriate function.  Parser
        # should have ensured that we don't have any unknown commands
        # here.
        self.args.handler()

        self.client.close()
        sys.exit(0)

    def cmd_info(self):
        """Print info about the server"""
        printf("Client library version: %s\n", self.client.client_version)
        printf("Server version: %s\n", self.client.version())
        printf("Server URL: %s\n", self.client.geturl())
        printf("Server database: %s\n", self.client.dbpath())

    def cmd_list(self):
        """List available streams"""
        streams = self.client.stream_list()
        for (path, layout) in streams:
            if (fnmatch.fnmatch(path, self.args.path) and
                fnmatch.fnmatch(layout, self.args.layout)):
                printf("%s %s\n", path, layout)

    def cmd_create(self):
        """Create new stream"""
        try:
            self.client.stream_create(self.args.path, self.args.layout)
        except nilmdb.client.ClientError as e:
            self.die("Error creating stream: %s\n", str(e))

    def cmd_metadata(self):
        """Manipulate metadata"""
        if self.args.set is not None or self.args.update is not None:
            # Either set, or update
            if self.args.set is not None:
                keyvals = self.args.set
                handler = self.client.stream_set_metadata
            else:
                keyvals = self.args.update
                handler = self.client.stream_update_metadata

            # Extract key=value pairs
            data = {}
            for keyval in keyvals:
                kv = keyval.split('=', 1)
                if len(kv) != 2 or kv[0] == "":
                    self.die("Error parsing key=value argument '%s'\n", keyval)
                data[kv[0]] = kv[1]

            # Make the call
            try:
                handler(self.args.path, data)
            except nilmdb.client.ClientError as e:
                self.die("Error setting/updating metadata: %s\n", str(e))
        else:
            # Get (or unspecified)
            keys = self.args.get or None
            try:
                data = self.client.stream_get_metadata(self.args.path, keys)
            except nilmdb.client.ClientError as e:
                self.die("Error getting metadata: %s\n", str(e))
            for key, value in sorted(data.items()):
                # Omit nonexistant keys
                if value is None:
                    value = ""
                printf("%s=%s\n", key, value)

    def cmd_insert(self):
        # Find requested stream
        streams = self.client.stream_list(self.args.path)
        if len(streams) != 1:
            self.die("Error getting stream info for path %s\n", self.args.path)

        layout = streams[0][1]

        if self.args.start and len(self.args.file) != 1:
            self.die("--start can only be used with one input file, for now")

        for filename in self.args.file:
            if filename == '-':
                process = None
                infile = sys.stdin
            else:
                if not os.path.exists(filename):
                    self.die("Error opening input file %s\n", filename)
                try:
                    # zcat is _much_ faster than python's gzopen
                    process = subprocess.Popen(["zcat", "-f", filename],
                                              bufsize = -1,
                                              stdin = None,
                                              stderr = None,
                                              stdout = PIPE)
                    infile = process.stdout
                except OSError: # pragma: no cover
                    self.die("Error spawning zcat process\n")

            # Build a timestamper for this file
            if self.args.none:
                ts = nilmdb.timestamper.TimestamperNull(infile)
            else:
                # If no rate, see if we can get it from nilmdb.layout
                if not self.args.rate:
                    try:
                        self.args.rate = nilmdb.layout.named[layout].rate_hz
                    except KeyError:
                        self.die("Need to specify --rate\n")

                # These will die if they can't parse
                if self.args.start:
                    start = self.parse_time(self.args.start)
                else:
                    start = self.parse_time(filename)

                ts = nilmdb.timestamper.TimestamperRate(infile, start, rate)

            print "Input file:", filename
            print "Timestamper:", ts
            print "Start:", start
            print "Rate:", rate

            self.die("not implemented")