nilmdb/nilmdb/cmdline/insert.py

import sys

from nilmdb.utils.printf import printf
import nilmdb.client
import nilmdb.utils.timestamper as timestamper
import nilmdb.utils.time


def setup(self, sub):
    cmd = sub.add_parser("insert", help="Insert data",
                         description="""
                         Insert data into a stream.
                         """)
    cmd.set_defaults(verify=cmd_insert_verify,
                     handler=cmd_insert)
    cmd.add_argument("-q", "--quiet", action='store_true',
                     help='suppress unnecessary messages')

    group = cmd.add_argument_group("Timestamping",
                                   description="""
                                   To add timestamps, specify the
                                   arguments --timestamp and --rate,
                                   and provide a starting time.
                                   """)

    group.add_argument("-t", "--timestamp", action="store_true",
                       help="Add timestamps to each line")
    group.add_argument("-r", "--rate", type=float,
                       help="Data rate, in Hz",
                       ).completer = self.complete.rate

    group = cmd.add_argument_group("Start time",
                                   description="""
                                   Start time may be manually
                                   specified with --start, or guessed
                                   from the filenames using
                                   --filename.  Set the TZ environment
                                   variable to change the default
                                   timezone.""")

    exc = group.add_mutually_exclusive_group()
    exc.add_argument("-s", "--start",
                     metavar="TIME", type=self.arg_time,
                     help="Starting timestamp (free-form)",
                     ).completer = self.complete.time
    exc.add_argument("-f", "--filename", action="store_true",
                     help="Use filename to determine start time")

    group = cmd.add_argument_group("End time",
                                   description="""
                                   End time for the overall stream.
                                   (required when not using --timestamp).
                                   Set the TZ environment
                                   variable to change the default
                                   timezone.""")
    group.add_argument("-e", "--end",
                       metavar="TIME", type=self.arg_time,
                       help="Ending timestamp (free-form)",
                       ).completer = self.complete.time

    group = cmd.add_argument_group("Required parameters")
    group.add_argument("path",
                       help="Path of stream, e.g. /foo/bar",
                       ).completer = self.complete.path
    group.add_argument("file", nargs='?', default='-',
                       help="File to insert (default: - (stdin))")
    return cmd


def cmd_insert_verify(self):
    if self.args.timestamp:
        if not self.args.rate:
            self.die("error: --rate is needed, but was not specified")
        if not self.args.filename and self.args.start is None:
            self.die("error: need --start or --filename "
                     "when adding timestamps")
    else:
        if self.args.start is None or self.args.end is None:
            self.die("error: when not adding timestamps, --start and "
                     "--end are required")


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", self.args.path)

    arg = self.args

    try:
        filename = arg.file
        if filename == '-':
            infile = sys.stdin.buffer
        else:
            try:
                infile = open(filename, "rb")
            except IOError:
                self.die("error opening input file %s", filename)

        if arg.start is None:
            try:
                arg.start = nilmdb.utils.time.parse_time(filename)
            except ValueError:
                self.die("error extracting start time from filename '%s'",
                         filename)

        if arg.timestamp:
            data = timestamper.TimestamperRate(infile, arg.start, arg.rate)
        else:
            data = iter(lambda: infile.read(1048576), b'')

        # Print info
        if not arg.quiet:
            printf(" Input file: %s\n", filename)
            printf(" Start time: %s\n",
                   nilmdb.utils.time.timestamp_to_human(arg.start))
            if arg.end:
                printf("   End time: %s\n",
                       nilmdb.utils.time.timestamp_to_human(arg.end))
            if arg.timestamp:
                printf("Timestamper: %s\n", str(data))

        # Insert the data
        self.client.stream_insert(arg.path, data, arg.start, arg.end)

    except nilmdb.client.Error as e:
        # TODO: It would be nice to be able to offer better errors
        # here, particularly in the case of overlap, which just shows
        # ugly bracketed ranges of 16-digit numbers and a mangled URL.
        # Need to consider adding something like e.prettyprint()
        # that is smarter about the contents of the error.
        self.die("error inserting data: %s", str(e))

    return