tableslurp

#!/usr/bin/env python
#
# -*- mode:python; sh-basic-offset:4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim:set tabstop=4 softtabstop=4 expandtab shiftwidth=4 fileencoding=utf-8:
#
# Copyright (c) 2012, Jorge A Gallegos <kad@blegh.net>
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import argparse
import boto
from dateutil import parser
import grp
import json
import logging
import threading
import os
import pwd
import socket
import sys
from Queue import Queue
import itertools

log = logging.getLogger('tableslurp')
stderr = logging.StreamHandler()
stderr.setFormatter(logging.Formatter(
    '%(name)s [%(asctime)s] %(levelname)s %(message)s'))
log.addHandler(stderr)
if os.environ.get('TDEBUG', False):
    log.setLevel(logging.DEBUG)
else:
    log.setLevel(logging.INFO)


class DownloadCounter(object):
    filename = None
    attemptcount = 0

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

    def increment(self):
        self.attemptcount += 1


class DownloadHandler(object):
    bucket_name = None
    owner = None
    group = None
    preserve = False
    target = None
    origin = None
    prefix = None
    force = False
    name = socket.getfqdn()
    fileset = []
    queue = Queue()
    num_threads = 4
    threads = {}

    def __init__(self, args):
        self.target = args.target[0]
        self.origin = args.origin[0]
        self.preserve = args.preserve
        self.bucket_name = args.bucket[0]
        self.num_threads = args.threads
        self.force = args.force
        if args.name:
            self.name = args.name
        self.prefix = '%s:%s' % (self.name, self.origin)

#       It may be a bit sub-optimal, but I rather fail sooner than later
        (owner, group) = self._build_file_set(args.file)

        if not self.preserve:
            owner = args.owner
            group = args.group

        try:
            self.owner = pwd.getpwnam(owner).pw_uid
            self.group = grp.getgrnam(group).gr_gid
        except Exception as e:
            log.error(e)
            raise OSError('User/Group pair %s:%s does not exist' %
                (owner, group,))

    def _get_bucket(self):
#       unsure if boto is thread-safe, will reconnect every time
        log.debug('Connecting to s3')
        conn = boto.connect_s3()
        bucket = conn.get_bucket(self.bucket_name, validate=False)
        log.debug('Connected to s3')
        return bucket

    def _build_file_set(self, target_file=None):
        log.info('Building fileset')
        keys = []
#       If you want to restore a file-set in particular
        bucket = self._get_bucket()
        if target_file:
            keys = [bucket.get_key('%s/%s-listdir.json' %
                (self.prefix, target_file))]
#       Otherwise try to fetch the most recent one
        else:
            recent_keys = dict()

            # Iterate over all keys in the bucket and save the most recently
            # modified file for each directory (Column Family). Assumes key
            # pattern like /path/to/your/<CF_NAME>/<FILENAME>
            for key in itertools.ifilter(
                    lambda k: k.key.endswith("-listdir.json"),
                    bucket.list(prefix='%s/' % (self.prefix,))):
                directory = key.key.split("/")[-2]
                if directory not in recent_keys:
                    recent_keys[directory] = key
                    continue

                if (parser.parse(recent_keys[directory].last_modified) <
                        parser.parse(key.last_modified)):
                    recent_keys[directory] = key

            keys = recent_keys.values()

        keys = list(itertools.ifilter(None, keys))

        if not keys:
            raise LookupError('Cannot find anything to restore from %s:%s/%s' %
                (bucket.name, self.prefix, target_file or ''))

        # Generate the fileset list by iterating over the keys for each CF and
        # retrieving the valid files from the JSON.
        for key in keys:
            json_data = json.loads(key.get_contents_as_string())
            assert len(json_data.keys()) == 1, "Invalid JSON format, aborting."

            # Prepend directory (CF name) to filename so we can restore to
            # that directory if we're restoring an entire keyspace.
            directory = json_data.keys()[0].split("/")[-1]
            filenames = ["%s/%s" % (directory, f) for f in json_data.values()[0]]
            self.fileset.extend(filenames)

        log.info('Fileset contains %d files to download' % (len(self.fileset)))
        k = bucket.get_key(keys[0].key)
#       The librato branch introduced this
        meta = k.get_metadata('stat')
        log.debug('Metadata is %s' % (meta,))
        owner = None
        group = None
        if meta:
            try:
                json_data = json.loads(meta)
                owner = json_data['user']
                group = json_data['group']
            except TypeError as te:
                log.debug(te)
                log.warning('Could not parse stat metadata for %s' % (k.name,))
            except KeyError as ke:
                log.debug(ke)
                log.warning('Incomplete stat metadata for %s, will ignore' %
                    (k.name,))
        return (owner, group)

    def _test_permissions(self):
        log.info('Will now try to test writing to the target dir %s' %
            (self.target,))
        try:

            if os.path.isdir(self.target) == False:
                log.debug('Creating temp file in %s' % (self.target,))
                os.makedirs(self.target)
            log.debug('Changing owner:group for %s to %s:%s' %
                (self.target, self.owner, self.group,))

            os.chown(self.target, self.owner, self.group)
        except Exception as e:
            log.debug(e)
            log.exception('%s exists' % (self.target,))
        log.info('Will write to %s' % (self.target,))

    def _worker(self, idx, queue):
        log.info('Thread #%d processing items' % (idx, ))
        bucket = self._get_bucket()

        while queue.empty() == False:
            queueddownload = queue.get()
            fname = queueddownload.filename
            directory = fname.split("/")[:1][0]
            # If the prefix ends with the name of a Column Family, we assume
            # that we are only restoring a single CF. If the prefix does not
            # end with the name of a CF we will create directories for each
            # individual CF.
            if self.prefix.endswith(directory):
                keypath = '%s/%s' % (self.prefix, fname.split("/")[1],)
                destfile = os.path.join(self.target, os.path.basename(fname))
            else:
                keypath = '%s/%s' % (self.prefix, fname,)
                destfile = os.path.join(self.target, directory, os.path.basename(fname))

            log.debug('Checking if we need to download %s to %s' %
                (keypath, destfile,))

            if queueddownload.attemptcount < 5:
                download = False
                #Retry downloading until we succeed
                try:
                    key = bucket.get_key(keypath)
                    log.debug('Key objectd is %s' % key)
                    if os.path.isfile(destfile):
                        stat = os.stat(destfile)
                        if self.force:
                            download = True
                        elif stat.st_size != key.size:
                            log.info('%s and %s size differs, will '
                                're-download' % (key.name, destfile,))
                            download = True
                    else:
                        download = True

                    if download:
                        log.info('Downloading %s from %s to %s' %
                            (key.name, bucket.name, destfile))
                        key.get_contents_to_filename(destfile)

                except Exception as e:
                    log.debug(e)
                    log.exception('Failed to download `%s` retrying' %
                        (fname,))
                    #We can't download, try again
                    queueddownload.increment()
                    queue.put(queueddownload)

            else:
                log.info('Tried to download %s too many times.  Giving up' %
                    fname)

            #Pop the task regardless of state.  If it fails we've put it back
            queue.task_done()

        log.info('Thread #%d finished processing' % (idx,))

    # If the prefix ends with the name of a Column Family, we assume that
    # we are only restoring a single CF. If the prefix does not end with
    # the name of a CF we create directories for each individual CF.
    def _create_cf_directories(self, cf_names):

        create_dirs = all(not self.prefix.endswith(cf) for cf in cf_names)

        if not create_dirs:
            return

        for cf_name in cf_names:
            directory = os.path.join(self.target, cf_name)
            try:
                os.mkdir(directory)
                os.chown(directory, self.owner, self.group)
                log.info('Created %s' % directory)
            except OSError, e:
                log.error('Error creating directory: %s' % e)
                pass

    def run(self):
        self._test_permissions()
        log.info('Running')
        cf_directories = set()

        #queue up the filesets
        for filename in self.fileset:
            log.info('Pushing file %s onto queue' % filename)
            self.queue.put(DownloadCounter(filename))
            cf_directories.add(os.path.split(filename)[0])

        self._create_cf_directories(cf_directories)

#       launch threads and attach an event to them
        for idx in range(0, self.num_threads):
            self.threads[idx] = {}
#            e = threading.Event()
            t = threading.Thread(target=self._worker,
                kwargs={'idx': idx, 'queue': self.queue})
            t.setDaemon(True)
            self.threads[idx] = t
            t.start()

        #Wait for everything to finish downloading
        self.queue.join()
        log.info('My job is done.')


def main():
    p = pwd.getpwnam(os.environ['USER'])
    owner = p.pw_name
    group = [_.gr_name for _ in grp.getgrall() if _.gr_gid == p.pw_gid][0]
    ap = argparse.ArgumentParser(
        description='This is the companion script to the `tablesnap` program '
        'which you can use to restore files from an Amazon S3 bucket to any '
        'given local directory which you have write-permissions on. While the '
        'code is straightforward, the program assumes the files you are '
        'restoring got previously backed up with `tablesnap`',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    ap.add_argument('-p', '--preserve', default=False, action='store_true',
        help='Preserve the permissions (if they exist) from the source. '
        'This overrides -o and -g')
    ap.add_argument('-o', '--owner', default=owner,
        help='After download, chown files to this user.')
    ap.add_argument('-g', '--group', default=group,
        help='After download, chgrp files to this group.')
    ap.add_argument('-t', '--threads', type=int, default=4,
        help='Split the download between this many threads')
    ap.add_argument('-f', '--file',
        help='If specified, will download the file-set this file belongs to '
        'instead of the latest one available.')
    ap.add_argument('--force', default=False, action='store_true',
        help='Force download files even if they exist')
    ap.add_argument('-n', '--name', default=socket.getfqdn(),
        help='Use this name instead of the FQDN to prefix the bucket dir')
    ap.add_argument('bucket', nargs=1,
        help='S3 bucket to download files from')
    ap.add_argument('origin', nargs=1,
        help='Path inside the bucket to the directory you want to download '
        'files from')
    ap.add_argument('target', nargs=1,
        help='Path in the local FS where files should be downloaded to')
    args = ap.parse_args()
    dh = DownloadHandler(args)
    dh.run()

if __name__ == '__main__':
    sys.exit(main())