def execute(self):
        logging.info('Ensuring the backup is found and is complete')
        if not self.cluster_backup.is_complete():
            raise Exception('Backup is not complete')

        # CASE 1 : We're restoring in place and a seed target has been provided
        if self.seed_target is not None:
            logging.info(
                'Restore will happen "In-Place", no new hardware is involved')
            self.in_place = True
            self.session_provider = CqlSessionProvider(
                [self.seed_target],
                username=self.config.cassandra.cql_username,
                password=self.config.cassandra.cql_password)

            with self.session_provider.new_session() as session:
                self._populate_ringmap(self.cluster_backup.tokenmap,
                                       session.tokenmap())

        # CASE 2 : We're restoring out of place, i.e. doing a restore test
        if self.host_list is not None:
            logging.info('Restore will happen on new hardware')
            self.in_place = False
            self._populate_hostmap()

        logging.info(
            'Starting Restore on all the nodes in this list: {}'.format(
                self.host_list))
        self._restore_data()
Beispiel #2
0
    def execute(self):
        logging.info('Ensuring the backup is found and is complete')
        if not self.cluster_backup.is_complete():
            raise RuntimeError('Backup is not complete')

        # CASE 1 : We're restoring using a seed target. Source/target mapping will be built based on tokenmap.
        if self.seed_target is not None:
            self.session_provider = CqlSessionProvider([self.seed_target],
                                                       self.config.cassandra)
            with self.session_provider.new_session() as session:
                self._populate_ringmap(self.cluster_backup.tokenmap,
                                       session.tokenmap())
                self._capture_release_version(session)

        # CASE 2 : We're restoring a backup on a different cluster
        if self.host_list is not None:
            logging.info('Restore will happen on new hardware')
            self.in_place = False
            self._populate_hostmap()
            self._capture_release_version(session=None)
            logging.info(
                'Starting Restore on all the nodes in this list: {}'.format(
                    self.host_list))

        self._restore_data()
def _get_cql_session_provider(config, hosts):

    if int(config.cassandra.is_ccm) == 1:
        cql_hosts = ['localhost']
    else:
        cql_hosts = hosts

    return CqlSessionProvider(cql_hosts, config.cassandra)
    def execute(self):
        # Two step: Take snapshot everywhere, then upload the backups to the external storage

        # Getting the list of Cassandra nodes.
        session_provider = CqlSessionProvider([self.config.storage.fqdn],
                                              self.config.cassandra)
        with session_provider.new_session() as session:
            tokenmap = session.tokenmap()
            self.hosts = [host for host in tokenmap.keys()]

        # First let's take a snapshot on all nodes at once
        # Here we will use parallelism of min(number of nodes, parallel_snapshots)
        logging.info('Creating snapshots on all nodes')
        self._create_snapshots()

        # Second
        logging.info('Uploading snapshots from nodes to external storage')
        self._upload_backup()
class RestoreJob(object):
    def __init__(self,
                 cluster_backup,
                 config,
                 temp_dir,
                 host_list,
                 seed_target,
                 keep_auth,
                 verify,
                 pssh_pool_size,
                 keyspaces={},
                 tables={},
                 bypass_checks=False,
                 use_sstableloader=False):
        self.id = uuid.uuid4()
        self.ringmap = None
        self.cluster_backup = cluster_backup
        self.session_provider = None
        self.config = config
        self.host_list = host_list
        self.seed_target = seed_target
        self.keep_auth = keep_auth
        self.verify = verify
        self.in_place = None
        self.temp_dir = temp_dir  # temporary files
        self.work_dir = self.temp_dir / 'medusa-job-{id}'.format(id=self.id)
        self.host_map = {
        }  # Map of backup host/target host for the restore process
        self.keyspaces = keyspaces
        self.tables = tables
        self.bypass_checks = bypass_checks
        self.use_sstableloader = use_sstableloader
        self.pssh_pool_size = pssh_pool_size
        self.cassandra = Cassandra(config.cassandra)
        fqdn_resolver = medusa.config.evaluate_boolean(
            self.config.cassandra.resolve_ip_addresses)
        self.fqdn_resolver = HostnameResolver(fqdn_resolver)

    def execute(self):
        logging.info('Ensuring the backup is found and is complete')
        if not self.cluster_backup.is_complete():
            raise Exception('Backup is not complete')

        # CASE 1 : We're restoring using a seed target. Source/target mapping will be built based on tokenmap.
        if self.seed_target is not None:
            self.session_provider = CqlSessionProvider([self.seed_target],
                                                       self.config.cassandra)

            with self.session_provider.new_session() as session:
                self._populate_ringmap(self.cluster_backup.tokenmap,
                                       session.tokenmap())

        # CASE 2 : We're restoring a backup on a different cluster
        if self.host_list is not None:
            logging.info('Restore will happen on new hardware')
            self.in_place = False
            self._populate_hostmap()
            logging.info(
                'Starting Restore on all the nodes in this list: {}'.format(
                    self.host_list))

        self._restore_data()

    def _pssh_run(self, hosts, command, hosts_variables=None):
        """
        Runs a command on hosts list using cstar under the hood
        There is no return made, to check the result there is a distinct function
        Return: True (success) or False (error)
        """
        logging.debug("Running pssh command on {}".format(hosts))
        pssh_run_success = False
        username = self.config.ssh.username if self.config.ssh.username != '' else None
        port = self.config.ssh.port
        pkey = None
        if self.config.ssh.key_file is not None and self.config.ssh.key_file != '':
            pkey = paramiko.RSAKey.from_private_key_file(
                self.config.ssh.key_file, None)

        client = ParallelSSHClient(hosts,
                                   forward_ssh_agent=True,
                                   pool_size=self.pssh_pool_size,
                                   user=username,
                                   port=port,
                                   pkey=pkey)
        logging.info('Executing "{}" on all nodes.'.format(command))
        output = client.run_command(command,
                                    host_args=hosts_variables,
                                    sudo=True)
        client.join(output)

        success = list(
            filter(
                lambda host_output: host_output.exit_code == 0,
                list(map(lambda host_output: host_output[1], output.items()))))
        error = list(
            filter(
                lambda host_output: host_output.exit_code != 0,
                list(map(lambda host_output: host_output[1], output.items()))))

        # Report on execution status
        if len(success) == len(hosts):
            logging.info(
                'Job executing "{}" ran and finished Successfully on all nodes.'
                .format(command))
            pssh_run_success = True
        elif len(error) > 0:
            logging.info(
                'Job executing "{}" ran and finished with errors on following nodes: {}'
                .format(
                    command,
                    sorted(
                        set(map(lambda host_output: host_output.host,
                                error)))))
            self.display_output(error)
        else:
            err_msg = 'Something unexpected happened while running pssh command'
            logging.error(err_msg)
            raise Exception(err_msg)

        return pssh_run_success

    def display_output(self, host_outputs):
        for host_out in host_outputs:
            for line in host_out.stdout:
                logging.info("{}-stdout: {}".format(host_out.host, line))
            for line in host_out.stderr:
                logging.info("{}-stderr: {}".format(host_out.host, line))

    def _validate_ringmap(self, tokenmap, target_tokenmap):
        for host, ring_item in target_tokenmap.items():
            if not ring_item.get('is_up'):
                raise Exception('Target {host} is not up!'.format(host=host))
        if len(target_tokenmap) != len(tokenmap):
            return False
        return True

    def _populate_ringmap(self, tokenmap, target_tokenmap):
        def _tokens_from_ringitem(ringitem):
            return ','.join(map(str, ringitem['tokens']))

        def _token_counts_per_host(tokenmap):
            for host, ringitem in tokenmap.items():
                return len(ringitem['tokens'])

        def _hosts_from_tokenmap(tokenmap):
            hosts = set()
            for host, ringitem in tokenmap.items():
                hosts.add(host)
            return hosts

        def _chunk(my_list, nb_chunks):
            groups = []
            for i in range(nb_chunks):
                groups.append([])
            for i in range(len(my_list)):
                groups[i % nb_chunks].append(my_list[i])
            return groups

        topology_matches = self._validate_ringmap(tokenmap, target_tokenmap)
        self.in_place = self._is_restore_in_place(tokenmap, target_tokenmap)
        if self.in_place:
            logging.info(
                "Restoring on the same cluster that was the backup was taken on (in place fashion)"
            )
            self.keep_auth = False
        else:
            logging.info(
                "Restoring on a different cluster than the backup one (remote fashion)"
            )
            if self.keep_auth:
                logging.info(
                    'system_auth keyspace will be left untouched on the target nodes'
                )
            else:
                # ops might not be aware of the underlying behavior towards auth. Let's ask what to do...
                really_keep_auth = None
                while (really_keep_auth != 'Y'
                       and really_keep_auth != 'n') and not self.bypass_checks:
                    really_keep_auth = input(
                        'Do you want to skip restoring the system_auth keyspace and keep the'
                        + ' credentials of the target cluster? (Y/n)')
                self.keep_auth = True if really_keep_auth == 'Y' else False

        if topology_matches:
            target_tokens = {
                _tokens_from_ringitem(ringitem): host
                for host, ringitem in target_tokenmap.items()
            }
            backup_tokens = {
                _tokens_from_ringitem(ringitem): host
                for host, ringitem in tokenmap.items()
            }

            target_tokens_per_host = _token_counts_per_host(tokenmap)
            backup_tokens_per_host = _token_counts_per_host(target_tokenmap)

            # we must have the same number of tokens per host in both vnode and normal clusters
            if target_tokens_per_host != backup_tokens_per_host:
                logging.info(
                    'Source/target rings have different number of tokens per node: {}/{}'
                    .format(backup_tokens_per_host, target_tokens_per_host))
                topology_matches = False

            # if not using vnodes, the tokens must match exactly
            if backup_tokens_per_host == 1 and target_tokens.keys(
            ) != backup_tokens.keys():
                extras = target_tokens.keys() ^ backup_tokens.keys()
                logging.info(
                    'Tokenmap is differently distributed. Extra items: {}'.
                    format(extras))
                topology_matches = False

        if topology_matches:
            # We can associate each restore node with exactly one backup node
            backup_ringmap = collections.defaultdict(list)
            target_ringmap = collections.defaultdict(list)
            for token, host in backup_tokens.items():
                backup_ringmap[token].append(host)
            for token, host in target_tokens.items():
                target_ringmap[token].append(host)

            self.ringmap = backup_ringmap
            i = 0
            for token, hosts in backup_ringmap.items():
                # take the node that has the same token list or pick the one with the same position in the map.
                restore_host = target_ringmap.get(
                    token,
                    list(target_ringmap.values())[i])[0]
                isSeed = True if self.fqdn_resolver.resolve_fqdn(
                    restore_host) in self._get_seeds_fqdn() else False
                self.host_map[restore_host] = {
                    'source': [hosts[0]],
                    'seed': isSeed
                }
                i += 1
        else:
            # Topologies are different between backup and restore clusters. Using the sstableloader for restore.
            self.use_sstableloader = True
            backup_hosts = _hosts_from_tokenmap(tokenmap)
            restore_hosts = list(_hosts_from_tokenmap(target_tokenmap))
            if len(backup_hosts) >= len(restore_hosts):
                grouped_backups = _chunk(list(backup_hosts),
                                         len(restore_hosts))
            else:
                grouped_backups = _chunk(list(backup_hosts), len(backup_hosts))
            for i in range(min([len(grouped_backups), len(restore_hosts)])):
                # associate one restore host with several backups as we don't have the same number of nodes.
                self.host_map[restore_hosts[i]] = {
                    'source': grouped_backups[i],
                    'seed': False
                }

    def _is_restore_in_place(self, backup_tokenmap, target_tokenmap):
        # If at least one node is part of both tokenmaps, then we're restoring in place
        # Otherwise we're restoring a remote cluster
        return len(set(backup_tokenmap.keys())
                   & set(target_tokenmap.keys())) > 0

    def _get_seeds_fqdn(self):
        seeds = list()
        for seed in self.cassandra.seeds:
            seeds.append(self.fqdn_resolver.resolve_fqdn(seed))
        return seeds

    def _populate_hostmap(self):
        with open(self.host_list, 'r') as f:
            for line in f.readlines():
                seed, target, source = line.replace('\n', '').split(
                    self.config.storage.host_file_separator)
                # in python, bool('False') evaluates to True. Need to test the membership as below
                self.host_map[self.fqdn_resolver.resolve_fqdn(target.strip())] \
                    = {'source': [self.fqdn_resolver.resolve_fqdn(source.strip())], 'seed': seed in ['True']}

    def _restore_data(self):
        # create workdir on each target host
        # Later: distribute a credential
        # construct command for each target host
        # invoke `nohup medusa-wrapper #{command}` on each target host
        # wait for exit on each
        logging.info('Starting cluster restore...')
        logging.info('Working directory for this execution: {}'.format(
            self.work_dir))
        for target, sources in self.host_map.items():
            logging.info(
                'About to restore on {} using {} as backup source'.format(
                    target, sources))

        logging.info(
            'This will delete all data on the target nodes and replace it with backup {}.'
            .format(self.cluster_backup.name))

        proceed = None
        while (proceed != 'Y' and proceed != 'n') and not self.bypass_checks:
            proceed = input('Are you sure you want to proceed? (Y/n)')

        if proceed == 'n':
            err_msg = 'Restore manually cancelled'
            logging.error(err_msg)
            raise Exception(err_msg)

        # work out which nodes are seeds in the target cluster
        target_seeds = [t for t, s in self.host_map.items() if s['seed']]
        logging.info("target seeds : {}".format(target_seeds))
        # work out which nodes are seeds in the target cluster
        target_hosts = self.host_map.keys()
        logging.info("target hosts : {}".format(target_hosts))

        if self.use_sstableloader is False:
            # stop all target nodes
            logging.info('Stopping Cassandra on all nodes currently up')

            # Generate a Job ID for this run
            job_id = str(uuid.uuid4())
            logging.debug('Job id is: {}'.format(job_id))
            # Define command to run
            command = self.config.cassandra.stop_cmd
            logging.debug('Command to run is: {}'.format(command))

            self._pssh_run(target_hosts, command, hosts_variables={})

        else:
            # we're using the sstableloader, which will require to (re)create the schema and empty the tables
            logging.info("Restoring schema on the target cluster")
            self._restore_schema()

        # trigger restores everywhere at once
        # pass in seed info so that non-seeds can wait for seeds before starting
        # seeds, naturally, don't wait for anything

        # Generate a Job ID for this run
        hosts_variables = []
        for target, source in [(t, s['source'])
                               for t, s in self.host_map.items()]:
            logging.info('Restoring data on {}...'.format(target))
            seeds = '' if target in target_seeds or len(target_seeds) == 0 \
                    else '--seeds {}'.format(','.join(target_seeds))
            hosts_variables.append((','.join(source), seeds))
            command = self._build_restore_cmd(target, source, seeds)

        pssh_run_success = self._pssh_run(target_hosts,
                                          command,
                                          hosts_variables=hosts_variables)

        if not pssh_run_success:
            # we could implement a retry.
            err_msg = 'Some nodes failed to restore. Exiting'
            logging.error(err_msg)
            raise Exception(err_msg)

        logging.info(
            'Restore process is complete. The cluster should be up shortly.')

        if self.verify:
            verify_restore(target_hosts, self.config)

    def _build_restore_cmd(self, target, source, seeds):
        in_place_option = '--in-place' if self.in_place else '--remote'
        keep_auth_option = '--keep-auth' if self.keep_auth else ''
        keyspace_options = expand_repeatable_option('keyspace', self.keyspaces)
        table_options = expand_repeatable_option('table', self.tables)
        # We explicitly set --no-verify since we are doing verification here in this module
        # from the control node
        verify_option = '--no-verify'

        # %s placeholders in the below command will get replaced by pssh using per host command substitution
        command = 'nohup sh -c "mkdir {work}; cd {work} && medusa-wrapper sudo medusa --fqdn=%s -vvv restore-node ' \
                  '{in_place} {keep_auth} %s {verify} --backup-name {backup} --temp-dir {temp_dir} ' \
                  '{use_sstableloader} {keyspaces} {tables}"' \
            .format(work=self.work_dir,
                    in_place=in_place_option,
                    keep_auth=keep_auth_option,
                    verify=verify_option,
                    backup=self.cluster_backup.name,
                    temp_dir=self.temp_dir,
                    use_sstableloader='--use-sstableloader' if self.use_sstableloader is True else '',
                    keyspaces=keyspace_options,
                    tables=table_options)

        logging.debug(
            'Restoring on node {} with the following command {}'.format(
                target, command))

        return command

    def _restore_schema(self):
        schema = parse_schema(self.cluster_backup.schema)
        with self.session_provider.new_session() as session:
            for keyspace in schema.keys():
                if keyspace.startswith("system"):
                    continue
                else:
                    self._create_or_recreate_schema_objects(
                        session, keyspace, schema[keyspace])

    def _create_or_recreate_schema_objects(self, session, keyspace,
                                           keyspace_schema):
        logging.info("(Re)creating schema for keyspace {}".format(keyspace))
        if (keyspace not in session.cluster.metadata.keyspaces):
            # Keyspace doesn't exist on the target cluster. Got to create it and all the tables as well.
            session.execute(keyspace_schema['create_statement'])
        for mv in keyspace_schema['materialized_views']:
            # MVs need to be dropped before we drop the tables
            logging.debug("Dropping MV {}.{}".format(keyspace, mv[0]))
            session.execute("DROP MATERIALIZED VIEW IF EXISTS {}.{}".format(
                keyspace, mv[0]))
        for table in keyspace_schema['tables'].items():
            logging.debug("Dropping table {}.{}".format(keyspace, table[0]))
            session.execute("DROP TABLE IF EXISTS {}.{}".format(
                keyspace, table[0]))
        for udt in keyspace_schema['udt'].items():
            # then custom types as they can be used in tables
            session.execute("DROP TYPE IF EXISTS {}.{}".format(
                keyspace, udt[0]))
            # Then we create the missing ones
            session.execute(udt[1])
        for table in keyspace_schema['tables'].items():
            logging.debug("Creating table {}.{}".format(keyspace, table[0]))
            # Create the tables
            session.execute(table[1])
        for index in keyspace_schema['indices'].items():
            # indices were dropped with their base tables
            logging.debug("Creating index {}.{}".format(keyspace, index[0]))
            session.execute(index[1])
        for mv in keyspace_schema['materialized_views']:
            # Base tables are created now, we can create the MVs
            logging.debug("Creating MV {}.{}".format(keyspace, mv[0]))
            session.execute(mv[1])
class RestoreJob(object):
    def __init__(self, cluster_backup, config, temp_dir, host_list, seed_target, keep_auth, verify,
                 keyspaces={}, tables={}, bypass_checks=False, use_sstableloader=False):
        self.id = uuid.uuid4()
        self.ringmap = None
        self.cluster_backup = cluster_backup
        self.session_provider = None
        self.config = config
        self.host_list = host_list
        self.seed_target = seed_target
        self.keep_auth = keep_auth
        self.verify = verify
        self.in_place = None
        self.temp_dir = temp_dir  # temporary files
        self.work_dir = self.temp_dir / 'medusa-job-{id}'.format(id=self.id)
        self.host_map = {}  # Map of backup host/target host for the restore process
        self.keyspaces = keyspaces
        self.tables = tables
        self.bypass_checks = bypass_checks
        self._ssh_agent_started = False
        self.use_sstableloader = use_sstableloader

    def execute(self):
        logging.info('Ensuring the backup is found and is complete')
        if not self.cluster_backup.is_complete():
            raise Exception('Backup is not complete')

        # CASE 1 : We're restoring in place and a seed target has been provided
        if self.seed_target is not None:
            logging.info('Restore will happen "In-Place", no new hardware is involved')
            self.in_place = True
            self.session_provider = CqlSessionProvider([self.seed_target],
                                                       username=self.config.cassandra.cql_username,
                                                       password=self.config.cassandra.cql_password)

            with self.session_provider.new_session() as session:
                self._populate_ringmap(self.cluster_backup.tokenmap, session.tokenmap())

        # CASE 2 : We're restoring out of place, i.e. doing a restore test
        if self.host_list is not None:
            logging.info('Restore will happen on new hardware')
            self.in_place = False
            self._populate_hostmap()

        logging.info('Starting Restore on all the nodes in this list: {}'.format(self.host_list))
        self._restore_data()
        if self._ssh_agent_started is True:
            self.ssh_cleanup()

    def _validate_ringmap(self, tokenmap, target_tokenmap):
        for host, ring_item in target_tokenmap.items():
            if not ring_item.get('is_up'):
                raise Exception('Target {host} is not up!'.format(host=host))
        if len(target_tokenmap) != len(tokenmap):
            return False
        return True

    def _populate_ringmap(self, tokenmap, target_tokenmap):

        def _tokens_from_ringitem(ringitem):
            return ','.join(map(str, ringitem['tokens']))

        def _token_counts_per_host(tokenmap):
            for host, ringitem in tokenmap.items():
                yield len(ringitem['tokens'])

        def _hosts_from_tokenmap(tokenmap):
            hosts = set()
            for host, ringitem in tokenmap.items():
                hosts.add(host)
            return hosts

        def _chunk(my_list, nb_chunks):
            groups = []
            for i in range(nb_chunks):
                groups.append([])
            for i in range(len(my_list)):
                groups[i % nb_chunks].append(my_list[i])
            return groups

        topology_matches = self._validate_ringmap(tokenmap, target_tokenmap)

        if topology_matches:
            target_tokens = {_tokens_from_ringitem(ringitem): host for host, ringitem in target_tokenmap.items()}
            backup_tokens = {_tokens_from_ringitem(ringitem): host for host, ringitem in tokenmap.items()}

            target_tokens_per_host = set(_token_counts_per_host(tokenmap))
            backup_tokens_per_host = set(_token_counts_per_host(target_tokenmap))

            # we must have the same number of tokens per host in both vnode and normal clusters
            if target_tokens_per_host != backup_tokens_per_host:
                logging.info('Source/target rings have different number of tokens per node: {}/{}'.format(
                    backup_tokens_per_host,
                    target_tokens_per_host
                ))
                topology_matches = False

            # if not using vnodes, the tokens must match exactly
            if len(backup_tokens_per_host) == 1 and target_tokens.keys() != backup_tokens.keys():
                extras = target_tokens.keys() ^ backup_tokens.keys()
                logging.info('Tokenmap is differently distributed. Extra items: {}'.format(extras))
                topology_matches = False

        if topology_matches:
            # We can associate each restore node with exactly one backup node
            ringmap = collections.defaultdict(list)
            for ring in backup_tokens, target_tokens:
                for token, host in ring.items():
                    ringmap[token].append(host)

            self.ringmap = ringmap
            for token, hosts in ringmap.items():
                self.host_map[hosts[1]] = {'source': [hosts[0]], 'seed': False}
        else:
            # Topologies are different between backup and restore clusters. Using the sstableloader for restore.
            self.use_sstableloader = True
            backup_hosts = _hosts_from_tokenmap(tokenmap)
            restore_hosts = list(_hosts_from_tokenmap(target_tokenmap))
            if len(backup_hosts) >= len(restore_hosts):
                grouped_backups = _chunk(list(backup_hosts), len(restore_hosts))
            else:
                grouped_backups = _chunk(list(backup_hosts), len(backup_hosts))
            for i in range(min([len(grouped_backups), len(restore_hosts)])):
                # associate one restore host with several backups as we don't have the same number of nodes.
                self.host_map[restore_hosts[i]] = {'source': grouped_backups[i], 'seed': False}

    def _populate_hostmap(self):
        with open(self.host_list, 'r') as f:
            for line in f.readlines():
                seed, target, source = line.replace('\n', '').split(self.config.storage.host_file_separator)
                # in python, bool('False') evaluates to True. Need to test the membership as below
                self.host_map[target.strip()] = {'source': [source.strip()], 'seed': seed in ['True']}

    def _restore_data(self):
        # create workdir on each target host
        # Later: distribute a credential
        # construct command for each target host
        # invoke `nohup medusa-wrapper #{command}` on each target host
        # wait for exit on each
        logging.info('Starting cluster restore...')
        logging.info('Working directory for this execution: {}'.format(self.work_dir))
        for target, sources in self.host_map.items():
            logging.info('About to restore on {} using {} as backup source'.format(target, sources))

        logging.info('This will delete all data on the target nodes and replace it with backup {}.'
                     .format(self.cluster_backup.name))

        proceed = None
        while (proceed != 'Y' and proceed != 'n') and not self.bypass_checks:
            proceed = input('Are you sure you want to proceed? (Y/n)')

        if proceed == 'n':
            err_msg = 'Restore manually cancelled'
            logging.error(err_msg)
            raise Exception(err_msg)

        if self.use_sstableloader is False:
            # stop all target nodes
            stop_remotes = []
            logging.info('Stopping Cassandra on all nodes')
            for target, source in [(t, s['source']) for t, s in self.host_map.items()]:
                client, connect_args = self._connect(target)
                if self.check_cassandra_running(target, client, connect_args):
                    logging.info('Cassandra is running on {}. Stopping it...'.format(target))
                    command = 'sh -c "{}"'.format(self.config.cassandra.stop_cmd)
                    stop_remotes.append(self._run(target, client, connect_args, command))
                else:
                    logging.info('Cassandra is not running on {}.'.format(target))

            # wait for all nodes to stop
            logging.info('Waiting for all nodes to stop...')
            finished, broken = self._wait_for(stop_remotes)
            if len(broken) > 0:
                err_msg = 'Some Cassandras failed to stop. Exiting'
                logging.error(err_msg)
                raise Exception(err_msg)
        else:
            # we're using the sstableloader, which will require to (re)create the schema and empty the tables
            logging.info("Restoring schema on the target cluster")
            self._restore_schema()

        # work out which nodes are seeds in the target cluster
        target_seeds = [t for t, s in self.host_map.items() if s['seed']]

        # trigger restores everywhere at once
        # pass in seed info so that non-seeds can wait for seeds before starting
        # seeds, naturally, don't wait for anything
        remotes = []
        for target, source in [(t, s['source']) for t, s in self.host_map.items()]:
            logging.info('Restoring data on {}...'.format(target))
            seeds = None if target in target_seeds else target_seeds
            remote = self._trigger_restore(target, source, seeds=seeds)
            remotes.append(remote)

        # wait for the restores
        logging.info('Starting to wait for the nodes to restore')
        finished, broken = self._wait_for(remotes)
        if len(broken) > 0:
            err_msg = 'Some nodes failed to restore. Exiting'
            logging.error(err_msg)
            raise Exception(err_msg)

        logging.info('Restore process is complete. The cluster should be up shortly.')

        if self.verify:
            hosts = list(map(lambda r: r.target, remotes))
            verify_restore(hosts, self.config)

    def _restore_schema(self):
        schema = parse_schema(self.cluster_backup.schema)
        with self.session_provider.new_session() as session:
            for keyspace in schema.keys():
                if keyspace.startswith("system"):
                    continue
                else:
                    self._create_or_recreate_schema_objects(session, keyspace, schema[keyspace])

    def _create_or_recreate_schema_objects(self, session, keyspace, keyspace_schema):
        logging.info("(Re)creating schema for keyspace {}".format(keyspace))
        if (keyspace not in session.cluster.metadata.keyspaces):
            # Keyspace doesn't exist on the target cluster. Got to create it and all the tables as well.
            session.execute(keyspace_schema['create_statement'])
        for mv in keyspace_schema['materialized_views']:
            # MVs need to be dropped before we drop the tables
            logging.debug("Dropping MV {}.{}".format(keyspace, mv[0]))
            session.execute("DROP MATERIALIZED VIEW {}.{}".format(keyspace, mv[0]))
        for table in keyspace_schema['tables'].items():
            logging.debug("Dropping table {}.{}".format(keyspace, table[0]))
            if table[0] in session.cluster.metadata.keyspaces[keyspace].tables.keys():
                # table already exists, drop it first
                session.execute("DROP TABLE {}.{}".format(keyspace, table[0]))
        for udt in keyspace_schema['udt'].items():
            # then custom types as they can be used in tables
            if udt[0] in session.cluster.metadata.keyspaces[keyspace].user_types.keys():
                # UDT already exists, drop it first
                session.execute("DROP TYPE {}.{}".format(keyspace, udt[0]))
            # Then we create the missing ones
            session.execute(udt[1])
        for table in keyspace_schema['tables'].items():
            logging.debug("Creating table {}.{}".format(keyspace, table[0]))
            # Create the tables
            session.execute(table[1])
        for index in keyspace_schema['indices'].items():
            # indices were dropped with their base tables
            logging.debug("Creating index {}.{}".format(keyspace, index[0]))
            session.execute(index[1])
        for mv in keyspace_schema['materialized_views']:
            # Base tables are created now, we can create the MVs
            logging.debug("Creating MV {}.{}".format(keyspace, mv[0]))
            session.execute(mv[1])

    def _trigger_restore(self, target, source, seeds=None):
        client, connect_args = self._connect(target)

        # TODO: If this command fails, the node is currently still marked as finished and not as broken.
        in_place_option = '--in-place' if self.in_place else ''
        keep_auth_option = '--keep-auth' if self.keep_auth else ''
        seeds_option = '--seeds {}'.format(','.join(seeds)) if seeds else ''
        keyspace_options = expand_repeatable_option('keyspace', self.keyspaces)
        table_options = expand_repeatable_option('table', self.tables)
        # We explicitly set --no-verify since we are doing verification here in this module
        # from the control node
        verify_option = '--no-verify'

        command = 'nohup sh -c "cd {work} && medusa-wrapper sudo medusa --fqdn={fqdn} -vvv restore-node ' \
                  '{in_place} {keep_auth} {seeds} {verify} --backup-name {backup} {use_sstableloader} ' \
                  '{keyspaces} {tables}"' \
            .format(work=self.work_dir,
                    fqdn=','.join(source),
                    in_place=in_place_option,
                    keep_auth=keep_auth_option,
                    seeds=seeds_option,
                    verify=verify_option,
                    backup=self.cluster_backup.name,
                    use_sstableloader='--use-sstableloader' if self.use_sstableloader is True else '',
                    keyspaces=keyspace_options,
                    tables=table_options)

        logging.debug('Restoring on node {} with the following command {}'.format(target, command))
        return self._run(target, client, connect_args, command)

    def _wait_for(self, remotes):
        finished, broken = [], []

        while True:
            time.sleep(5)  # TODO: configure sleep

            if len(remotes) == len(finished) + len(broken):
                # TODO: make a nicer exit condition
                logging.info('Exiting because all jobs are done.')
                break

            for i, remote in enumerate(remotes):

                if remote in broken or remote in finished:
                    continue

                # If the remote does not set an exit status and the channel closes
                # the exit_status is negative.
                logging.debug('remote.channel.exit_status: {}'.format(remote.channel.exit_status))
                if remote.channel.exit_status_ready and remote.channel.exit_status >= 0:
                    if remote.channel.exit_status == 0:
                        finished.append(remote)
                        logging.info('Command succeeded on {}'.format(remote.target))
                    else:
                        broken.append(remote)
                        logging.error('Command failed on {} : '.format(remote.target))
                        logging.error('Output : {}'.format(remote.stdout.readlines()))
                        logging.error('Err output : {}'.format(remote.stderr.readlines()))
                        try:
                            stderr = self.read_file(remote, self.work_dir / 'stderr')
                        except IOError:
                            stderr = 'There was no stderr file'
                        logging.error(stderr)
                    # We got an exit code that does not indicate an error, but not necessarily
                    # success. Cleanup channel and move to next remote.
                    remote.channel.close()
                    # also close the client. this will free file descriptors
                    # in case we start re-using remotes this close will need to go away
                    remote.client.close()
                    continue

                if remote.client.get_transport().is_alive() and not remote.channel.closed:
                    # Send an ignored packet for keep alive and later noticing a broken connection
                    logging.debug('Keeping {} alive.'.format(remote.target))
                    remote.client.get_transport().send_ignore()
                else:
                    client = paramiko.client.SSHClient()
                    client.load_system_host_keys()
                    client.connect(**remote.connect_args)

                    # TODO: check pid to exist before assuming medusa-wrapper to pick it up
                    command = 'cd {work}; medusa-wrapper'.format(work=self.work_dir)
                    remotes[i] = self._run(remote.target, client, remote.connect_args, command)

        if len(broken) > 0:
            logging.info('Command failed on the following nodes:')
            for remote in broken:
                logging.info(remote.target)
        else:
            logging.info('Commands succeeded on all nodes')

        return finished, broken

    def _connect(self, target):
        logging.debug('Connecting to {}'.format(target))

        pkey = None
        if self.config.ssh.key_file is not None and self.config.ssh.key_file != '':
            pkey = paramiko.RSAKey.from_private_key_file(self.config.ssh.key_file, None)
            if self._ssh_agent_started is False:
                self.create_agent()
                add_key_cmd = '{} {}'.format(SSH_ADD_KEYS_CMD, self.config.ssh.key_file)
                subprocess.check_output(add_key_cmd, universal_newlines=True, shell=True)
                self._ssh_agent_started = True

        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.client.AutoAddPolicy())
        connect_args = {
            'hostname': target,
            'username': self.config.ssh.username,
            'pkey': pkey,
            'compress': True,
            'password': None
        }
        client.connect(**connect_args)

        logging.debug('Successfully connected to {}'.format(target))
        sftp = client.open_sftp()
        try:
            sftp.mkdir(str(self.work_dir))
        except OSError:
            err_msg = 'Working directory {} on {} failed.' \
                      'Folder might exist already, ignoring exception'.format(str(self.work_dir), target)
            logging.debug(err_msg)
        except Exception as ex:
            err_msg = 'Creating working directory on {} failed: {}'.format(target, str(ex))
            logging.error(err_msg)
            raise Exception(err_msg)
        finally:
            sftp.close()

        return client, connect_args

    def _run(self, target, client, connect_args, command):
        transport = client.get_transport()
        session = transport.open_session()
        session.get_pty()
        paramiko.agent.AgentRequestHandler(session)
        session.exec_command(command.replace('sudo', 'sudo -S'))
        bufsize = -1
        stdout = session.makefile('r', bufsize)
        stderr = session.makefile_stderr('r', bufsize)
        logging.debug('Running \'{}\' remotely on {}'.format(command, connect_args['hostname']))
        return Remote(target, connect_args, client, stdout.channel, stdout, stderr)

    def read_file(self, remote, remotepath):
        with remote.client.open_sftp() as ftp_client:
            with ftp_client.file(remotepath.as_posix(), 'r') as f:
                return str(f.read(), 'utf-8')

    def check_cassandra_running(self, host, client, connect_args):
        command = 'sh -c "{}"'.format(self.config.cassandra.check_running)
        remote = self._run(host, client, connect_args, command)
        return remote.channel.recv_exit_status() == 0

    def create_agent(self):
        """
        Function that creates the agent and sets the environment variables.
        """
        output = subprocess.check_output(SSH_AGENT_CREATE_CMD, universal_newlines=True, shell=True)
        if output:
            output = output.strip().split('\n')
            for item in output[0:2]:
                envvar, val = item.split(';')[0].split('=')
                logging.debug('Setting environment variable: {}={}'.format(envvar, val))
                os.environ[envvar] = val

    def ssh_cleanup(self):
        """
        Function that kills the agents created so that there aren't too many agents lying around eating up resources.
        """
        # Kill the agent
        subprocess.check_output(SSH_AGENT_KILL_CMD, universal_newlines=True, shell=True)
        # Reset these values so that other function
        os.environ[SSH_AUTH_SOCK_ENVVAR] = ''
        os.environ[SSH_AGENT_PID_ENVVAR] = ''
class RestoreJob(object):
    def __init__(self,
                 cluster_backup,
                 config,
                 temp_dir,
                 host_list,
                 seed_target,
                 keep_auth,
                 verify,
                 parallel_restores,
                 keyspaces=None,
                 tables=None,
                 bypass_checks=False,
                 use_sstableloader=False,
                 version_target=None):

        self.id = uuid.uuid4()
        self.ringmap = None
        self.cluster_backup = cluster_backup
        self.session_provider = None
        self.orchestration = Orchestration(config, parallel_restores)
        self.config = config
        self.host_list = host_list
        self.seed_target = seed_target
        self.keep_auth = keep_auth
        self.verify = verify
        self.in_place = None
        self.temp_dir = temp_dir  # temporary files
        self.work_dir = self.temp_dir / 'medusa-job-{id}'.format(id=self.id)
        self.host_map = {
        }  # Map of backup host/target host for the restore process
        self.keyspaces = keyspaces if keyspaces else {}
        self.tables = tables if tables else {}
        self.bypass_checks = bypass_checks
        self.use_sstableloader = use_sstableloader
        self.pssh_pool_size = parallel_restores
        self.cassandra = Cassandra(config)
        fqdn_resolver = medusa.utils.evaluate_boolean(
            self.config.cassandra.resolve_ip_addresses)
        self.fqdn_resolver = HostnameResolver(fqdn_resolver)
        self._version_target = version_target

    def execute(self):
        logging.info('Ensuring the backup is found and is complete')
        if not self.cluster_backup.is_complete():
            raise RuntimeError('Backup is not complete')

        # CASE 1 : We're restoring using a seed target. Source/target mapping will be built based on tokenmap.
        if self.seed_target is not None:
            self.session_provider = CqlSessionProvider([self.seed_target],
                                                       self.config.cassandra)
            with self.session_provider.new_session() as session:
                self._populate_ringmap(self.cluster_backup.tokenmap,
                                       session.tokenmap())
                self._capture_release_version(session)

        # CASE 2 : We're restoring a backup on a different cluster
        if self.host_list is not None:
            logging.info('Restore will happen on new hardware')
            self.in_place = False
            self._populate_hostmap()
            self._capture_release_version(session=None)
            logging.info(
                'Starting Restore on all the nodes in this list: {}'.format(
                    self.host_list))

        self._restore_data()

    @staticmethod
    def _validate_ringmap(tokenmap, target_tokenmap):
        for host, ring_item in target_tokenmap.items():
            if not ring_item.get('is_up'):
                raise RuntimeError(
                    'Target {host} is not up!'.format(host=host))
        if len(target_tokenmap) != len(tokenmap):
            return False
        return True

    def _populate_ringmap(self, tokenmap, target_tokenmap):
        def _tokens_from_ringitem(ringitem):
            return ','.join(map(str, ringitem['tokens']))

        def _token_counts_per_host(tokenmap):
            for host, ringitem in tokenmap.items():
                return len(ringitem['tokens'])

        def _hosts_from_tokenmap(tokenmap):
            hosts = set()
            for host, ringitem in tokenmap.items():
                hosts.add(host)
            return hosts

        def _chunk(my_list, nb_chunks):
            groups = []
            for i in range(nb_chunks):
                groups.append([])
            for i in range(len(my_list)):
                groups[i % nb_chunks].append(my_list[i])
            return groups

        target_tokens = {}
        backup_tokens = {}
        topology_matches = self._validate_ringmap(tokenmap, target_tokenmap)
        self.in_place = self._is_restore_in_place(tokenmap, target_tokenmap)
        if self.in_place:
            logging.info(
                "Restoring on the same cluster that was the backup was taken on (in place fashion)"
            )
            self.keep_auth = False
        else:
            logging.info(
                "Restoring on a different cluster than the backup one (remote fashion)"
            )
            if self.keep_auth:
                logging.info(
                    'system_auth keyspace will be left untouched on the target nodes'
                )
            else:
                # ops might not be aware of the underlying behavior towards auth. Let's ask what to do...
                really_keep_auth = None
                while (really_keep_auth != 'Y'
                       and really_keep_auth != 'n') and not self.bypass_checks:
                    really_keep_auth = input(
                        'Do you want to skip restoring the system_auth keyspace and keep the'
                        + ' credentials of the target cluster? (Y/n)')
                self.keep_auth = True if really_keep_auth == 'Y' else False

        if topology_matches:
            target_tokens = {
                _tokens_from_ringitem(ringitem): host
                for host, ringitem in target_tokenmap.items()
            }
            backup_tokens = {
                _tokens_from_ringitem(ringitem): host
                for host, ringitem in tokenmap.items()
            }

            target_tokens_per_host = _token_counts_per_host(tokenmap)
            backup_tokens_per_host = _token_counts_per_host(target_tokenmap)

            # we must have the same number of tokens per host in both vnode and normal clusters
            if target_tokens_per_host != backup_tokens_per_host:
                logging.info(
                    'Source/target rings have different number of tokens per node: {}/{}'
                    .format(backup_tokens_per_host, target_tokens_per_host))
                topology_matches = False

            # if not using vnodes, the tokens must match exactly
            if backup_tokens_per_host == 1 and target_tokens.keys(
            ) != backup_tokens.keys():
                extras = target_tokens.keys() ^ backup_tokens.keys()
                logging.info(
                    'Tokenmap is differently distributed. Extra items: {}'.
                    format(extras))
                topology_matches = False

        if topology_matches:
            # We can associate each restore node with exactly one backup node
            backup_ringmap = collections.defaultdict(list)
            target_ringmap = collections.defaultdict(list)
            for token, host in backup_tokens.items():
                backup_ringmap[token].append(host)
            for token, host in target_tokens.items():
                target_ringmap[token].append(host)

            self.ringmap = backup_ringmap
            i = 0
            for token, hosts in backup_ringmap.items():
                # take the node that has the same token list or pick the one with the same position in the map.
                restore_host = target_ringmap.get(
                    token,
                    list(target_ringmap.values())[i])[0]
                is_seed = True if self.fqdn_resolver.resolve_fqdn(
                    restore_host) in self._get_seeds_fqdn() else False
                self.host_map[restore_host] = {
                    'source': [hosts[0]],
                    'seed': is_seed
                }
                i += 1
            logging.debug("self.host_map: {}".format(self.host_map))
        else:
            # Topologies are different between backup and restore clusters. Using the sstableloader for restore.
            self.use_sstableloader = True
            backup_hosts = _hosts_from_tokenmap(tokenmap)
            restore_hosts = list(_hosts_from_tokenmap(target_tokenmap))
            if len(backup_hosts) >= len(restore_hosts):
                grouped_backups = _chunk(list(backup_hosts),
                                         len(restore_hosts))
            else:
                grouped_backups = _chunk(list(backup_hosts), len(backup_hosts))
            for i in range(min([len(grouped_backups), len(restore_hosts)])):
                # associate one restore host with several backups as we don't have the same number of nodes.
                self.host_map[restore_hosts[i]] = {
                    'source': grouped_backups[i],
                    'seed': False
                }

    @staticmethod
    def _is_restore_in_place(backup_tokenmap, target_tokenmap):
        # If at least one node is part of both tokenmaps, then we're restoring in place
        # Otherwise we're restoring a remote cluster
        return len(set(backup_tokenmap.keys())
                   & set(target_tokenmap.keys())) > 0

    def _get_seeds_fqdn(self):
        seeds = list()
        for seed in self.cassandra.seeds:
            seeds.append(self.fqdn_resolver.resolve_fqdn(seed))
        logging.debug("seeds are: {}".format(seeds))
        return seeds

    def _populate_hostmap(self):
        """
        When there are no seed nodes to pull cluster topology from, the essential information required for a restore
            can be passed in via a simple file using the --host-list CLI argument.

        Each line in the file must have three pieces of information in this order:
            - the string `True` or `False`; This indicates if the source node was a seed node
            - the host/ip that the restore operation is to take place on / destination node
            - the host/ip where the data came from / source node
        Each field is separated by a comma.

        E.G.: Using medusa to restore a 4 node cluster from a previous backup taken of that same cluster:
            medusa@cassandra-node01:~$ cat nodes.list
            True,10.10.1.127,10.10.1.127
            True,10.10.1.128,10.10.1.128
            False,10.10.1.129,10.10.1.129
            False,10.10.1.130,10.10.1.130

        :return:
        """
        with open(self.host_list, 'r') as f:
            for line in f.readlines():
                # Remove leading/trailing whitespace
                _line = line.strip()
                # Ignore comment lines
                if _line.startswith('#'):
                    continue
                seed, target, source = _line.split(
                    self.config.storage.host_file_separator)
                # in python, bool('False') evaluates to True. Need to test the membership as below
                target_resolved = self.fqdn_resolver.resolve_fqdn(
                    target.strip())
                source_resolved = self.fqdn_resolver.resolve_fqdn(
                    source.strip())
                self.host_map[target_resolved] = {
                    'source': [source_resolved],
                    'seed': seed in ['True']
                }

    def _restore_data(self):
        # create workdir on each target host
        # Later: distribute a credential
        # construct command for each target host
        # invoke `nohup medusa-wrapper #{command}` on each target host
        # wait for exit on each
        logging.info('Starting cluster restore...')
        logging.info('Working directory for this execution: {}'.format(
            self.work_dir))
        for target, sources in self.host_map.items():
            logging.info(
                'About to restore on {} using {} as backup source'.format(
                    target, sources))

        logging.info(
            "This will delete all data on the target nodes and replace it with backup '{}'."
            .format(self.cluster_backup.name))

        proceed = None
        while (proceed != 'Y' and proceed != 'n') and not self.bypass_checks:
            proceed = input('Are you sure you want to proceed? (Y/n)')

        if proceed == 'n':
            err_msg = 'Restore manually cancelled'
            logging.error(err_msg)
            raise RuntimeError(err_msg)

        # work out which nodes are seeds in the target cluster
        target_seeds = [t for t, s in self.host_map.items() if s['seed']]
        logging.info("target seeds : {}".format(target_seeds))
        # work out which nodes are seeds in the target cluster
        target_hosts = [host for host in self.host_map.keys()]
        logging.info("target hosts : {}".format(target_hosts))

        if self.use_sstableloader is False:
            # stop all target nodes
            logging.info('Stopping Cassandra on all nodes currently up')

            # Generate a Job ID for this run
            job_id = str(uuid.uuid4())
            logging.debug('Job id is: {}'.format(job_id))
            # Define command to run
            command = self.config.cassandra.stop_cmd
            logging.debug('Command to run is: {}'.format(command))

            self.orchestration.pssh_run(target_hosts,
                                        command,
                                        hosts_variables={})

        else:
            # we're using the sstableloader, which will require to (re)create the schema and empty the tables
            logging.info("Restoring schema on the target cluster")
            self._restore_schema()

        # trigger restores everywhere at once
        # pass in seed info so that non-seeds can wait for seeds before starting
        # seeds, naturally, don't wait for anything

        # Generate a Job ID for this run
        hosts_variables = []
        for target, source in [(t, s['source'])
                               for t, s in self.host_map.items()]:
            logging.info('Restoring data on {}...'.format(target))
            seeds = '' if target in target_seeds or len(target_seeds) == 0 \
                else '--seeds {}'.format(','.join(target_seeds))
            hosts_variables.append((','.join(source), seeds))

        command = self._build_restore_cmd()
        pssh_run_success = self.orchestration.pssh_run(
            target_hosts, command, hosts_variables=hosts_variables)

        if not pssh_run_success:
            # we could implement a retry.
            err_msg = 'Some nodes failed to restore. Exiting'
            logging.error(err_msg)
            raise RuntimeError(err_msg)

        logging.info(
            'Restore process is complete. The cluster should be up shortly.')

        if self.verify:
            verify_restore(target_hosts, self.config)

    def _build_restore_cmd(self):
        in_place_option = '--in-place' if self.in_place else '--remote'
        keep_auth_option = '--keep-auth' if self.keep_auth else ''
        keyspace_options = expand_repeatable_option('keyspace', self.keyspaces)
        table_options = expand_repeatable_option('table', self.tables)
        # We explicitly set --no-verify since we are doing verification here in this module
        # from the control node
        verify_option = '--no-verify'

        # %s placeholders in the below command will get replaced by pssh using per host command substitution
        command = 'mkdir -p {work}; cd {work} && medusa-wrapper {sudo} medusa {config} ' \
                  '--fqdn=%s -vvv restore-node ' \
                  '{in_place} {keep_auth} %s {verify} --backup-name {backup} --temp-dir {temp_dir} ' \
                  '{use_sstableloader} {keyspaces} {tables}' \
            .format(work=self.work_dir,
                    sudo='sudo' if medusa.utils.evaluate_boolean(self.config.cassandra.use_sudo) else '',
                    config=f'--config-file {self.config.file_path}' if self.config.file_path else '',
                    in_place=in_place_option,
                    keep_auth=keep_auth_option,
                    verify=verify_option,
                    backup=self.cluster_backup.name,
                    temp_dir=self.temp_dir,
                    use_sstableloader='--use-sstableloader' if self.use_sstableloader else '',
                    keyspaces=keyspace_options,
                    tables=table_options)

        logging.debug(
            'Preparing to restore on all nodes with the following command: {}'.
            format(command))

        return command

    def _restore_schema(self):
        schema = parse_schema(self.cluster_backup.schema)
        with self.session_provider.new_session() as session:
            for keyspace in schema.keys():
                if keyspace.startswith("system"):
                    continue
                else:
                    self._create_or_recreate_schema_objects(
                        session, keyspace, schema[keyspace])

    def _create_or_recreate_schema_objects(self, session, keyspace,
                                           keyspace_schema):
        logging.info("(Re)creating schema for keyspace {}".format(keyspace))
        if keyspace not in session.cluster.metadata.keyspaces:
            # Keyspace doesn't exist on the target cluster. Got to create it and all the tables as well.
            session.execute(keyspace_schema['create_statement'])
        for mv in keyspace_schema['materialized_views']:
            # MVs need to be dropped before we drop the tables
            logging.debug("Dropping MV {}.{}".format(keyspace, mv[0]))
            session.execute("DROP MATERIALIZED VIEW IF EXISTS {}.{}".format(
                keyspace, mv[0]))
        for table in keyspace_schema['tables'].items():
            logging.debug("Dropping table {}.{}".format(keyspace, table[0]))
            session.execute("DROP TABLE IF EXISTS {}.{}".format(
                keyspace, table[0]))
        for udt in keyspace_schema['udt'].items():
            # then custom types as they can be used in tables
            session.execute("DROP TYPE IF EXISTS {}.{}".format(
                keyspace, udt[0]))
            # Then we create the missing ones
            session.execute(udt[1])
        for table in keyspace_schema['tables'].items():
            logging.debug("Creating table {}.{}".format(keyspace, table[0]))
            # Create the tables
            session.execute(table[1])
        for index in keyspace_schema['indices'].items():
            # indices were dropped with their base tables
            logging.debug("Creating index {}.{}".format(keyspace, index[0]))
            session.execute(index[1])
        for mv in keyspace_schema['materialized_views']:
            # Base tables are created now, we can create the MVs
            logging.debug("Creating MV {}.{}".format(keyspace, mv[0]))
            session.execute(mv[1])

    # Capture release version as specified, from driver, or use default.
    # This is necessary for logic that requires knowledge of differences between 2, 3, and 4.
    def _capture_release_version(self, session):
        # If no version specified via CLI, but have a session, get version from driver.
        if not self._version_target and session:
            driver_app_version = session.cluster.application_version
            if driver_app_version:
                logging.debug('Driver version provided as: {}'.format(
                    driver_app_version))
                HostMan.set_release_version(driver_app_version)
            else:
                logging.debug(
                    'Unable to obtain app_version via driver or command line, '
                    'using default: {}'.format(
                        HostMan.DEFAULT_RELEASE_VERSION))
                # Using default as target wasn't found by driver or provided to RestoreJob
                HostMan.set_release_version(HostMan.DEFAULT_RELEASE_VERSION)
        # If no session available or specified version from CLI, use default.
        elif not self._version_target:
            # Use default
            HostMan.set_release_version(HostMan.DEFAULT_RELEASE_VERSION)
        else:
            # Use what is specified from CLI as version.
            logging.debug('Target version provided as: {}'.format(
                self._version_target))
            HostMan.set_release_version(self._version_target)