<commit_before>import re
import shutil
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from datetime import datetime
from fabric.api import env
from fabric.api import execute
from fabric.api import hide
from fabric.api import put
from fabric.api import sudo
from fabric.context_managers import settings
from multiprocessing.dummy import Pool
import json
import logging
import os
from tempfile import NamedTemporaryFile
import time
import sys


class Snapshot(object):
    """
    A Snapshot instance keeps the details about a cassandra snapshot

    Multiple snaphosts can be stored in a single S3 bucket

    A Snapshot is best described by:
        - its name (which defaults to the utc time of creation)
        - the list of hostnames the snapshot runs on
        - the list of keyspaces being backed up
        - the keyspace table being backed up
        - the S3 bucket's base path where the snapshot is stored

    Snapshots data (and incremental backups) are stored using the following convention:

        s3_bucket_name:/<base_path>/<snapshot_name>/<node-hostname>/...

    Snapshots are represented on S3 by their manifest file, this makes incremental backups
    much easier
    """

    SNAPSHOT_TIMESTAMP_FORMAT = '%Y%m%d%H%M%S'

    def __init__(self, base_path, s3_bucket, hosts, keyspaces, table):
        self.s3_bucket = s3_bucket
        self.name = self.make_snapshot_name()
        self.hosts = hosts
        self.keyspaces = keyspaces
        self.table = table
        self._base_path = base_path

    def dump_manifest_file(self):
        manifest_data = {
            'name': self.name,
            'base_path': self._base_path,
            'hosts': self.hosts,
            'keyspaces': self.keyspaces,
            'table': self.table
        }
        return json.dumps(manifest_data)

    @staticmethod
    def load_manifest_file(data, s3_bucket):
        manifest_data = json.loads(data)
        snapshot = Snapshot(
            base_path=manifest_data['base_path'],
            s3_bucket=s3_bucket,
            hosts=manifest_data['hosts'],
            keyspaces=manifest_data['keyspaces'],
            table=manifest_data['table']
        )
        snapshot.name = manifest_data['name']
        return snapshot

    @property
    def base_path(self):
        return '/'.join([self._base_path, self.name])

    def make_snapshot_name(self):
        return datetime.utcnow().strftime(self.SNAPSHOT_TIMESTAMP_FORMAT)

    def unix_time_name(self):
        dt = datetime.strptime(self.name, self.SNAPSHOT_TIMESTAMP_FORMAT)
        return time.mktime(dt.timetuple()) * 1000

    def __cmp__(self, other):
        return self.unix_time_name() - other.unix_time_name()

    def __repr__(self):
        return self.name

    __str__ = __repr__


class RestoreWorker(object):
    def __init__(self, aws_access_key_id, aws_secret_access_key, snapshot):
        self.aws_secret_access_key = aws_secret_access_key
        self.aws_access_key_id = aws_access_key_id
        self.s3connection = S3Connection(aws_access_key_id=self.aws_access_key_id,
                                         aws_secret_access_key=self.aws_secret_access_key)
        self.snapshot = snapshot
        self.keyspace_table_matcher = None

    def restore(self, keyspace, table, hosts, target_hosts):
        # TODO:
        # 4. sstableloader

        logging.info("Restoring keyspace=%(keyspace)s, table=%(table)s" % dict(keyspace=keyspace,
                                                                               table=table))

        logging.info("From hosts: %(hosts)s to: %(target_hosts)s" % dict(hosts=', '.join(hosts),
                                                                         target_hosts=', '.join(
                                                                             target_hosts)))
        if not table:
            table = ".*?"

        bucket = self.s3connection.get_bucket(self.snapshot.s3_bucket)

        matcher_string = "(%(hosts)s).*/(%(keyspace)s)/(%(table)s)" % dict(hosts='|'.join(hosts), keyspace=keyspace, table=table)
        self.keyspace_table_matcher = re.compile(matcher_string)

        keys = []
        tables = set()

        for k in bucket.list(self.snapshot.base_path):
            r = self.keyspace_table_matcher.search(k.name)
            if not r:
                continue

            tables.add(r.group(3))
            keys.append(k)

        self._delete_old_dir_and_create_new(keyspace, tables)

        total_size = reduce(lambda s, k: s + k.size, keys, 0)

        logging.info("Found %(files_count)d files, with total size of %(size)d." % dict(
            files_count=len(keys),
            size=total_size))

        self._download_keys(keys, total_size)

        logging.info("Finished downloading...")

        self._run_sstableloader(keyspace, tables, target_hosts)

    def _delete_old_dir_and_create_new(self, keyspace, tables):
        keyspace_path = './%s' % keyspace
        if os.path.exists(keyspace_path) and os.path.isdir(keyspace_path):
            logging.warning("Deleteing directory (%s)..." % keyspace_path)
            shutil.rmtree(keyspace_path)

        for table in tables:
            path = './%s/%s' % (keyspace, table)
            if not os.path.exists(path):
                os.makedirs(path)

    def _download_keys(self, keys, total_size, pool_size=5):
        logging.info("Starting to download...")

        progress_string = ""
        read_bytes = 0

        thread_pool = Pool(pool_size)

        for size in thread_pool.imap(self._download_key, keys):
            old_width = len(progress_string)
            read_bytes += size
            progress_string = "%d / %d bytes (%.2f%%)" % (read_bytes, total_size, (read_bytes/float(total_size))*100.0)
            width = len(progress_string)
            padding = ""
            if width < old_width:
                padding = " "*(width-old_width)
            progress_string = "%s%s\r" % (progress_string, padding)

            sys.stderr.write(progress_string)

    def _download_key(self, key):
        r = self.keyspace_table_matcher.search(key.name)
        filename = './%s/%s/%s_%s' % (r.group(2), r.group(3), key.name.split('/')[2], key.name.split('/')[-1])
        key.get_contents_to_filename(filename)

        return key.size

    def _run_sstableloader(self, keyspace, tables, target_hosts):
        # TODO: get path to sstableloader
        for table in tables:
            command = 'sstableloader --nodes %(hosts)s -v %(keyspace)s/%(table)s' % dict(
                hosts=target_hosts, keyspace=keyspace, table=table)
            logging.info("invoking: %s", command)

            os.system(command)

class BackupWorker(object):
    """
    BackupWorker does the actual snapshot / backup work

    Backup process is split in this steps:
        - requesting cassandra to create new backups
        - uploading backup files to S3
        - clearing backup files from nodes
        - updating backup meta information

    When performing a new snapshot the manifest of the snapshot is
    uploaded to S3 for later use.

    Snapshot's manifest path:
    /<snapshot_base_path>/<snapshot_name>/manifest.json

    Everytime a backup is done a description of the current ring is
    saved next to the snapshot manifest file

    """
<commit_msg>Comma separate the target hosts<commit_after>import re
import shutil
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from datetime import datetime
from fabric.api import env
from fabric.api import execute
from fabric.api import hide
from fabric.api import put
from fabric.api import sudo
from fabric.context_managers import settings
from multiprocessing.dummy import Pool
import json
import logging
import os
from tempfile import NamedTemporaryFile
import time
import sys


class Snapshot(object):
    """
    A Snapshot instance keeps the details about a cassandra snapshot

    Multiple snaphosts can be stored in a single S3 bucket

    A Snapshot is best described by:
        - its name (which defaults to the utc time of creation)
        - the list of hostnames the snapshot runs on
        - the list of keyspaces being backed up
        - the keyspace table being backed up
        - the S3 bucket's base path where the snapshot is stored

    Snapshots data (and incremental backups) are stored using the following convention:

        s3_bucket_name:/<base_path>/<snapshot_name>/<node-hostname>/...

    Snapshots are represented on S3 by their manifest file, this makes incremental backups
    much easier
    """

    SNAPSHOT_TIMESTAMP_FORMAT = '%Y%m%d%H%M%S'

    def __init__(self, base_path, s3_bucket, hosts, keyspaces, table):
        self.s3_bucket = s3_bucket
        self.name = self.make_snapshot_name()
        self.hosts = hosts
        self.keyspaces = keyspaces
        self.table = table
        self._base_path = base_path

    def dump_manifest_file(self):
        manifest_data = {
            'name': self.name,
            'base_path': self._base_path,
            'hosts': self.hosts,
            'keyspaces': self.keyspaces,
            'table': self.table
        }
        return json.dumps(manifest_data)

    @staticmethod
    def load_manifest_file(data, s3_bucket):
        manifest_data = json.loads(data)
        snapshot = Snapshot(
            base_path=manifest_data['base_path'],
            s3_bucket=s3_bucket,
            hosts=manifest_data['hosts'],
            keyspaces=manifest_data['keyspaces'],
            table=manifest_data['table']
        )
        snapshot.name = manifest_data['name']
        return snapshot

    @property
    def base_path(self):
        return '/'.join([self._base_path, self.name])

    def make_snapshot_name(self):
        return datetime.utcnow().strftime(self.SNAPSHOT_TIMESTAMP_FORMAT)

    def unix_time_name(self):
        dt = datetime.strptime(self.name, self.SNAPSHOT_TIMESTAMP_FORMAT)
        return time.mktime(dt.timetuple()) * 1000

    def __cmp__(self, other):
        return self.unix_time_name() - other.unix_time_name()

    def __repr__(self):
        return self.name

    __str__ = __repr__


class RestoreWorker(object):
    def __init__(self, aws_access_key_id, aws_secret_access_key, snapshot):
        self.aws_secret_access_key = aws_secret_access_key
        self.aws_access_key_id = aws_access_key_id
        self.s3connection = S3Connection(aws_access_key_id=self.aws_access_key_id,
                                         aws_secret_access_key=self.aws_secret_access_key)
        self.snapshot = snapshot
        self.keyspace_table_matcher = None

    def restore(self, keyspace, table, hosts, target_hosts):
        # TODO:
        # 4. sstableloader

        logging.info("Restoring keyspace=%(keyspace)s, table=%(table)s" % dict(keyspace=keyspace,
                                                                               table=table))

        logging.info("From hosts: %(hosts)s to: %(target_hosts)s" % dict(hosts=', '.join(hosts),
                                                                         target_hosts=', '.join(
                                                                             target_hosts)))
        if not table:
            table = ".*?"

        bucket = self.s3connection.get_bucket(self.snapshot.s3_bucket)

        matcher_string = "(%(hosts)s).*/(%(keyspace)s)/(%(table)s)" % dict(hosts='|'.join(hosts), keyspace=keyspace, table=table)
        self.keyspace_table_matcher = re.compile(matcher_string)

        keys = []
        tables = set()

        for k in bucket.list(self.snapshot.base_path):
            r = self.keyspace_table_matcher.search(k.name)
            if not r:
                continue

            tables.add(r.group(3))
            keys.append(k)

        self._delete_old_dir_and_create_new(keyspace, tables)

        total_size = reduce(lambda s, k: s + k.size, keys, 0)

        logging.info("Found %(files_count)d files, with total size of %(size)d." % dict(
            files_count=len(keys),
            size=total_size))

        self._download_keys(keys, total_size)

        logging.info("Finished downloading...")

        self._run_sstableloader(keyspace, tables, target_hosts)

    def _delete_old_dir_and_create_new(self, keyspace, tables):
        keyspace_path = './%s' % keyspace
        if os.path.exists(keyspace_path) and os.path.isdir(keyspace_path):
            logging.warning("Deleteing directory (%s)..." % keyspace_path)
            shutil.rmtree(keyspace_path)

        for table in tables:
            path = './%s/%s' % (keyspace, table)
            if not os.path.exists(path):
                os.makedirs(path)

    def _download_keys(self, keys, total_size, pool_size=5):
        logging.info("Starting to download...")

        progress_string = ""
        read_bytes = 0

        thread_pool = Pool(pool_size)

        for size in thread_pool.imap(self._download_key, keys):
            old_width = len(progress_string)
            read_bytes += size
            progress_string = "%d / %d bytes (%.2f%%)" % (read_bytes, total_size, (read_bytes/float(total_size))*100.0)
            width = len(progress_string)
            padding = ""
            if width < old_width:
                padding = " "*(width-old_width)
            progress_string = "%s%s\r" % (progress_string, padding)

            sys.stderr.write(progress_string)

    def _download_key(self, key):
        r = self.keyspace_table_matcher.search(key.name)
        filename = './%s/%s/%s_%s' % (r.group(2), r.group(3), key.name.split('/')[2], key.name.split('/')[-1])
        key.get_contents_to_filename(filename)

        return key.size

    def _run_sstableloader(self, keyspace, tables, target_hosts):
        # TODO: get path to sstableloader
        for table in tables:
            command = 'sstableloader --nodes %(hosts)s -v %(keyspace)s/%(table)s' % dict(
                hosts=','.join(target_hosts), keyspace=keyspace, table=table)
            logging.info("invoking: %s", command)

            os.system(command)

class BackupWorker(object):
    """
    BackupWorker does the actual snapshot / backup work

    Backup process is split in this steps:
        - requesting cassandra to create new backups
        - uploading backup files to S3
        - clearing backup files from nodes
        - updating backup meta information

    When performing a new snapshot the manifest of the snapshot is
    uploaded to S3 for later use.

    Snapshot's manifest path:
    /<snapshot_base_path>/<snapshot_name>/manifest.json

    Everytime a backup is done a description of the current ring is
    saved next to the snapshot manifest file

    """

{'fraction_non_alphanumeric': 0.0966952264381885, 'fraction_numerical': 0.005031959744322046, 'mean_word_length': 2.325870646766169, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 34, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4609842', 'n_tokens_mistral': 4390, 'n_tokens_neox': 4219, 'n_words': 1052}