Esempio n. 1
0
class FTSFileOperation(FileTransferOperation, FileTransferQuery, FileDeletionOperation, FileDeletionQuery):
    def __init__(self, config):
        FileTransferOperation.__init__(self, config)
        FileTransferQuery.__init__(self, config)
        FileDeletionOperation.__init__(self, config)
        FileDeletionQuery.__init__(self, config)

        self.server_url = config.fts_server
        self.server_id = 0 # server id in the DB

        # Parameter "retry" for fts3.new_job. 0 = server default
        self.fts_retry = config.get('fts_retry', 0)

        # String passed to fts3.new_*_job(metadata = _)
        self.metadata_string = config.get('metadata_string', 'Dynamo')

        # Proxy to be forwarded to FTS
        self.x509proxy = config.get('x509proxy', None)
        self.x509proxy_orig = config.get('x509proxy', None)

        # Bookkeeping device
        self.db = MySQL(config.db_params)

        # Reuse the context object
        self.keep_context = config.get('keep_context', False)
        self._context = None

    def num_pending_transfers(self): #override
        # Check the number of files in queue
        # We first thought about counting files with /files, but FTS seems to return only 1000 maximum even when "limit" is set much larger
        #files = self._ftscallurl('/files?state_in=ACTIVE,SUBMITTED,READY&limit=%d' % self.max_pending_transfers)
        #return len(files)

        num_pending = 0
        file_states = ['SUBMITTED', 'READY', 'ACTIVE', 'STAGING', 'STARTED']

        jobs = self._ftscall('list_jobs', state_in = ['SUBMITTED', 'ACTIVE', 'STAGING'])
        from random import shuffle
        shuffle(jobs)

        total_count = 0
        for job in jobs:
            total_count = total_count + 1
            #LOG.info("List_files call 1")
            job_info = self._ftscall('get_job_status', job['job_id'], list_files = True)
            for file_info in job_info['files']:
                if file_info['file_state'] in file_states:
                    num_pending += 1
                    if num_pending == self.max_pending_transfers + 1:
                        # don't need to query more
                        return num_pending

        return num_pending

    def num_pending_deletions(self): #override
        # See above
        #files = self._ftscallurl('/files?state_in=ACTIVE,SUBMITTED,READY&limit=%d' % self.max_pending_deletions)
        #return len(files)

        num_pending = 0
        file_states = ['SUBMITTED', 'READY', 'ACTIVE']

        jobs = self._ftscall('list_jobs', state_in = ['SUBMITTED', 'ACTIVE'])
        from random import shuffle
        shuffle(jobs)

        total_count = 0
        for job in jobs:
            total_count = total_count + 1
            #LOG.info("List_files call 2")
            job_info = self._ftscall('get_job_status', job['job_id'], list_files = True)
            for file_info in job_info['dm']:
                if file_info['file_state'] in file_states:
                    num_pending += 1
                    if num_pending == self.max_pending_deletions + 1:
                        # don't need to query more
                        return num_pending

        return num_pending

    def form_batches(self, tasks): #override
        if len(tasks) == 0:
            return []

        # FTS3 has no restriction on how to group the transfers, but cannot apparently take thousands
        # of tasks at once
        batches = [[]]
        for task in tasks:
            batches[-1].append(task)
            if len(batches[-1]) == self.batch_size:
                batches.append([])

        return batches

    def start_transfers(self, batch_id, batch_tasks): #override
        result = {}

        stage_files = []
        transfers = []

        s_pfn_to_task = {}
        t_pfn_to_task = {}

        for task in batch_tasks:
            sub = task.subscription
            lfn = sub.file.lfn
            dest_pfn = sub.destination.to_pfn(lfn, 'gfal2')
            source_pfn = task.source.to_pfn(lfn, 'gfal2')

            self.x509proxy = sub.destination.x509proxy

            if task.source.storage_type == Site.TYPE_MSS:
                self.x509proxy = task.source.x509proxy

            if dest_pfn is None or source_pfn is None:
                # either gfal2 is not supported or lfn could not be mapped
                LOG.warning('Could not obtain PFN for %s at %s or %s', lfn, sub.destination.name, task.source.name)
                result[task] = False
                continue

            if self.checksum_algorithm:
                checksum = '%s:%s' % (self.checksum_algorithm, str(sub.file.checksum[self.checksum_index]))
                verify_checksum = 'target'
            else:
                checksum = None
                verify_checksum = False

            if task.source.storage_type == Site.TYPE_MSS:
                LOG.debug('Staging %s at %s', lfn, task.source.name)

                # need to stage first
                stage_files.append((source_pfn, dest_pfn, checksum, sub.file.size))

                # task identified by the source PFN
                s_pfn_to_task[source_pfn] = task
            else:
                LOG.info("Here we are")
                LOG.info('Submitting transfer of %s from %s to %s to FTS', lfn, task.source.name, sub.destination.name)

                transfers.append(fts3.new_transfer(source_pfn, dest_pfn, checksum = checksum, filesize = sub.file.size))

                # there should be only one task per destination pfn
                t_pfn_to_task[dest_pfn] = task

        if len(stage_files) != 0:
            LOG.debug('Submit new staging job for %d files', len(stage_files))
            job = fts3.new_staging_job([ff[0] for ff in stage_files], bring_online = 36000, metadata = self.metadata_string)
            success = self._submit_job(job, 'staging', batch_id, dict((pfn, task.id) for pfn, task in s_pfn_to_task.iteritems()), x509=task.source.x509proxy)

            for source_pfn, _, _, _ in stage_files:
                result[s_pfn_to_task[source_pfn]] = success

            if success and not self._read_only:
                LOG.debug('Recording staging queue')
                fields = ('id', 'source', 'destination', 'checksum', 'size')
                mapping = lambda ff: (s_pfn_to_task[ff[0]].id,) + ff
                if not self._read_only:
                    self.db.insert_many('fts_staging_queue', fields, mapping, stage_files)

        if len(transfers) != 0:
            LOG.debug('Submit new transfer job for %d files', len(transfers))
            LOG.info("Submitting transfer job from disk to site %s with proxy %s." % (sub.destination.name, sub.destination.x509proxy))

            job = fts3.new_job(transfers, retry = self.fts_retry, overwrite = True, 
                               verify_checksum = verify_checksum, metadata = self.metadata_string)
            success = self._submit_job(job, 'transfer', batch_id, dict((pfn, task.id) for pfn, task in t_pfn_to_task.iteritems()), x509=sub.destination.x509proxy)

            for transfer in transfers:
                dest_pfn = transfer['destinations'][0]
                result[t_pfn_to_task[dest_pfn]] = success

        return result

    def start_deletions(self, batch_id, batch_tasks): #override
        result = {}

        pfn_to_task = {}

        for task in batch_tasks:
            desub = task.desubscription
            lfn = desub.file.lfn
            pfn = desub.site.to_pfn(lfn, 'gfal2')

            if pfn is None:
                # either gfal2 is not supported or lfn could not be mapped
                result[task] = False
                continue

            # there should be only one task per destination pfn
            pfn_to_task[pfn] = task

        job = fts3.new_delete_job(pfn_to_task.keys(), metadata = self.metadata_string)

        success = self._submit_job(job, 'deletion', batch_id, dict((pfn, task.id) for pfn, task in pfn_to_task.iteritems()))

        for task in pfn_to_task.itervalues():
            result[task] = success

        return result

    def cancel_transfers(self, task_ids): #override
        return self._cancel(task_ids, 'transfer')

    def cancel_deletions(self, task_ids): #override
        return self._cancel(task_ids, 'deletion')

    def cleanup(self): #override
        sql = 'DELETE FROM f USING `fts_transfer_tasks` AS f'
        sql += ' LEFT JOIN `transfer_tasks` AS t ON t.`id` = f.`id`'
        sql += ' LEFT JOIN `fts_transfer_batches` AS b ON b.`id` = f.`fts_batch_id`'
        sql += ' WHERE t.`id` IS NULL OR b.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM f USING `fts_staging_queue` AS f'
        sql += ' LEFT JOIN `fts_transfer_tasks` AS t ON t.`id` = f.`id`'
        sql += ' WHERE t.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM f USING `fts_deletion_tasks` AS f'
        sql += ' LEFT JOIN `deletion_tasks` AS t ON t.`id` = f.`id`'
        sql += ' LEFT JOIN `fts_deletion_batches` AS b ON b.`id` = f.`fts_batch_id`'
        sql += ' WHERE t.`id` IS NULL OR b.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM f USING `fts_transfer_batches` AS f'
        sql += ' LEFT JOIN `transfer_batches` AS t ON t.`id` = f.`batch_id`'
        sql += ' WHERE t.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM f USING `fts_deletion_batches` AS f'
        sql += ' LEFT JOIN `deletion_batches` AS t ON t.`id` = f.`batch_id`'
        sql += ' WHERE t.`id` IS NULL'
        self.db.query(sql)

        # Delete the source tasks - caution: wipes out all tasks when switching the operation backend
        sql = 'DELETE FROM t USING `transfer_tasks` AS t'
        sql += ' LEFT JOIN `fts_transfer_tasks` AS f ON f.`id` = t.`id`'
        sql += ' WHERE f.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM t USING `deletion_tasks` AS t'
        sql += ' LEFT JOIN `fts_deletion_tasks` AS f ON f.`id` = t.`id`'
        sql += ' WHERE f.`id` IS NULL'
        self.db.query(sql)

    def get_transfer_status(self, batch_id): #override
        if self.server_id == 0:
            self._set_server_id()

        results = self._get_status(batch_id, 'transfer')

        staged_tasks = []

        for task_id, status, exitcode, msg, start_time, finish_time in self._get_status(batch_id, 'staging'):
            if status == FileQuery.STAT_DONE:
                staged_tasks.append(task_id)
                results.append((task_id, FileQuery.STAT_QUEUED, -1, None, None, None))
            else:
                # these tasks won't appear in results from _get_status('transfer')
                # because no transfer jobs have been submitted yet
                results.append((task_id, status, exitcode, None, start_time, finish_time))

        if len(staged_tasks) != 0:
            if self.checksum_algorithm:
                verify_checksum = 'target'
            else:
                verify_checksum = None

            which_dest = ""
            transfers = []
            pfn_to_tid = {}

            for task_id, source_pfn, dest_pfn, checksum, filesize in self.db.select_many('fts_staging_queue', ('id', 'source', 'destination', 'checksum', 'size'), 'id', staged_tasks):
                for s in inventory.sites.itervalues():
                    if s.backend in dest_pfn and s.backend != "":
                        which_dest = s
                #transfers = []
                #pfn_to_tid = {}
                transfers.append(fts3.new_transfer(source_pfn, dest_pfn, checksum = checksum, filesize = filesize))
                pfn_to_tid[dest_pfn] = task_id

            job = fts3.new_job(transfers, retry = self.fts_retry, overwrite = True, verify_checksum = verify_checksum, metadata = self.metadata_string)
                
                #which_dest = ""
                #for s in inventory.sites.itervalues():
                #    if s.backend in dest_pfn and s.backend != "":
                #        which_dest = s
            LOG.info("Submitting job from buffer to site %s with proxy %s." % (which_dest.name, which_dest.x509proxy))
            #LOG.info("dest_pfn: %s" % dest_pfn) 
            success = self._submit_job(job, 'transfer', batch_id, pfn_to_tid, x509=which_dest.x509proxy)

            if success and not self._read_only:
                self.db.delete_many('fts_staging_queue', 'id', pfn_to_tid.values())

        return results

    def get_deletion_status(self, batch_id): #override
        if self.server_id == 0:
            self._set_server_id()

        return self._get_status(batch_id, 'deletion')

    def write_transfer_history(self, history_db, task_id, history_id): #override
        self._write_history(history_db, task_id, history_id, 'transfer')

    def write_deletion_history(self, history_db, task_id, history_id): #override
        self._write_history(history_db, task_id, history_id, 'deletion')

    def forget_transfer_status(self, task_id): #override
        return self._forget_status(task_id, 'transfer')

    def forget_deletion_status(self, task_id): #override
        return self._forget_status(task_id, 'deletion')

    def forget_transfer_batch(self, task_id): #override
        return self._forget_batch(task_id, 'transfer')

    def forget_deletion_batch(self, task_id): #override
        return self._forget_batch(task_id, 'deletion')

    def _ftscall(self, method, *args, **kwd):
        return self._do_ftscall(binding = (method, args, kwd))

    def _ftscallurl(self, url):
        # Call to FTS URLs that don't have python bindings
        return self._do_ftscall(url = url)

    def _do_ftscall(self, binding = None, url = None):
        
        proxy = self.x509proxy

        if binding is not None:
            method, args, kwd = binding
            for k in kwd.iteritems():
                LOG.info(k)
                if k[0] == "ucert":
                    LOG.info("Setting proxy to %s" % k[1])
                    proxy = k[1]

        if self._context is None:
            # request_class = Request -> use "requests"-based https call (instead of default PyCURL,
            # which may not be able to handle proxy certificates depending on the cURL installation)
            # verify = False -> do not verify the server certificate
            LOG.info("Setting context with proxy %s" % proxy)
            context = fts3.Context(self.server_url, ucert = proxy, ukey = proxy,
                                   request_class = Request, verify = False)

            if self.keep_context:
                self._context = context
        else:
            context = self._context

        if binding is not None:
            reqstring = binding[0]
        else:
            reqstring = url

        LOG.debug('FTS: %s', reqstring)

        wait_time = 1.

        for attempt in xrange(10):
            try:
                if binding is not None:
                    method, args, kwd = binding

                    if "ucert" in kwd:
                        del kwd["ucert"]
                    if "ukey" in kwd:
                        del kwd["ukey"]

                    return getattr(fts3, method)(context, *args, **kwd)
                else:
                    return json.loads(context.get(url))
            except fts_exceptions.ServerError as exc:
                if str(exc.reason) == '500':
                    # Internal server error - let's try again
                    pass
            except fts_exceptions.TryAgain:
                pass

            time.sleep(wait_time)
            wait_time *= 1.5

        LOG.error('Failed to communicate with FTS server: %s', reqstring)
        raise RuntimeError('Failed to communicate with FTS server: %s' % reqstring)

    def _submit_job(self, job, optype, batch_id, pfn_to_tid, x509=None):
        if self._read_only:
            job_id = 'test'
        else:
            try:
                if x509 is not None:
                    job_id = self._ftscall('submit', job, ucert=x509, ukey=x509)
                else:
                    job_id = self._ftscall('submit', job)
            except:
                exc_type, exc, tb = sys.exc_info()
                LOG.error('Failed to submit %s to FTS: Exception %s (%s)', optype, exc_type.__name__, str(exc))
                return False

        # list of file-level operations (one-to-one with pfn)
        try:
            if optype == 'transfer' or optype == 'staging':
                key = 'files'
            else:
                key = 'dm'

            #LOG.info("List_files call 3")
            fts_files = self._ftscall('get_job_status', job_id = job_id, list_files = True)[key]
        except:
            exc_type, exc, tb = sys.exc_info()
            LOG.error('Failed to get status of job %s from FTS: Exception %s (%s)', job_id, exc_type.__name__, str(exc))
            return False

        if self.server_id == 0:
            self._set_server_id()

        if optype == 'transfer' or optype == 'staging':
            table_name = 'fts_transfer_batches'
            columns = ('batch_id', 'task_type', 'fts_server_id', 'job_id')
            values = (batch_id, optype, self.server_id, job_id)
        else:
            table_name = 'fts_deletion_batches'
            columns = ('batch_id', 'fts_server_id', 'job_id')
            values = (batch_id, self.server_id, job_id)

        if not self._read_only:
            fts_batch_id = self.db.insert_get_id(table_name, columns = columns, values = values)

        if optype == 'transfer' or optype == 'staging':
            table_name = 'fts_transfer_tasks'
            pfn_key = 'dest_surl'
        else:
            table_name = 'fts_deletion_tasks'
            pfn_key = 'source_surl'

        fields = ('id', 'fts_batch_id', 'fts_file_id')
        mapping = lambda f: (pfn_to_tid[f[pfn_key]], fts_batch_id, f['file_id'])

        if not self._read_only:
            self.db.insert_many(table_name, fields, mapping, fts_files, do_update = True, update_columns = ('fts_batch_id', 'fts_file_id'))

        return True

    def _cancel(self, task_ids, optype):
        sql = 'SELECT b.`job_id`, f.`fts_file_id` FROM `fts_{op}_tasks` AS f'
        sql += ' INNER JOIN `fts_{op}_batches` AS b ON b.`id` = f.`fts_batch_id`'
        result = self.db.execute_many(sql.format(op = optype), MySQL.bare('f.`id`'), task_ids)

        by_job = collections.defaultdict(list)

        for job_id, file_id in result:
            by_job[job_id].append(file_id)

        if not self._read_only:
            for job_id, ids in by_job.iteritems():
                try:
                    self._ftscall('cancel', job_id, file_ids = ids)
                except:
                    LOG.error('Failed to cancel FTS job %s', job_id)
    
    def _get_status(self, batch_id, optype):
        if optype == 'transfer' or optype == 'staging':
            sql = 'SELECT `id`, `job_id` FROM `fts_transfer_batches`'
            sql += ' WHERE `task_type` = %s AND `fts_server_id` = %s AND `batch_id` = %s'
            batch_data = self.db.query(sql, optype, self.server_id, batch_id)
            task_table_name = 'fts_transfer_tasks'
        else:
            sql = 'SELECT `id`, `job_id` FROM `fts_deletion_batches`'
            sql += ' WHERE `fts_server_id` = %s AND `batch_id` = %s'
            batch_data = self.db.query(sql, self.server_id, batch_id)
            task_table_name = 'fts_deletion_tasks'

        message_pattern = re.compile('(?:DESTINATION|SOURCE|TRANSFER|DELETION) \[([0-9]+)\] (.*)')

        results = []

        for fts_batch_id, job_id in batch_data:
            LOG.debug('Checking status of FTS %s batch %s', optype, job_id)

            sql = 'SELECT `fts_file_id`, `id` FROM `{table}` WHERE `fts_batch_id` = %s'.format(table = task_table_name)
            fts_to_task = dict(self.db.xquery(sql, fts_batch_id))

            try:
                #LOG.info("List_files call 4")
                result = self._ftscall('get_job_status', job_id = job_id, list_files = True)
            except:
                LOG.error('Failed to get job status for FTS job %s', job_id)
                LOG.error(optype)
                continue
    
            if optype == 'transfer' or optype == 'staging':
                fts_files = result['files']
            else:
                fts_files = result['dm']

            for fts_file in fts_files:
                try:
                    task_id = fts_to_task[fts_file['file_id']]
                except KeyError:
                    continue
    
                state = fts_file['file_state']
                exitcode = -1
                start_time = None
                finish_time = None
                get_time = False

                try:
                    message = fts_file['reason']
                except KeyError:
                    message = None

                if message is not None:
                    # Check if reason follows a known format (from which we can get the exit code)
                    matches = message_pattern.match(message)
                    if matches is not None:
                        exitcode = int(matches.group(1))
                        message = matches.group(2)
                    # Additionally, if the message is a known one, convert the exit code
                    c = find_msg_code(message)
                    if c is not None:
                        exitcode = c

                    # HDFS site with gridftp-hdfs gives a I/O error (500) when the file is not there
                    if optype == 'deletion' and 'Input/output error' in message:
                        exitcode = errno.ENOENT

                if state == 'FINISHED':
                    status = FileQuery.STAT_DONE
                    exitcode = 0
                    get_time = True

                elif state == 'FAILED':
                    status = FileQuery.STAT_FAILED
                    get_time = True

                elif state == 'CANCELED':
                    status = FileQuery.STAT_CANCELLED
                    get_time = True

                elif state == 'SUBMITTED':
                    status = FileQuery.STAT_NEW

                else:
                    status = FileQuery.STAT_QUEUED

                if optype == 'transfer' and exitcode == errno.EEXIST:
                    # Transfer + destination exists -> not an error
                    status = FileQuery.STAT_DONE
                    exitcode = 0
                elif optype == 'deletion' and exitcode == errno.ENOENT:
                    # Deletion + destination does not exist -> not an error
                    status = FileQuery.STAT_DONE
                    exitcode = 0
                    
                if get_time:
                    try:
                        start_time = calendar.timegm(time.strptime(fts_file['start_time'], '%Y-%m-%dT%H:%M:%S'))
                    except TypeError: # start time is NULL (can happen when the job is cancelled)
                        start_time = None
                    try:
                        finish_time = calendar.timegm(time.strptime(fts_file['finish_time'], '%Y-%m-%dT%H:%M:%S'))
                    except TypeError:
                        start_time = None

                LOG.debug('%s %d: %s, %d, %s, %s, %s', optype, task_id, FileQuery.status_name(status), exitcode, message, start_time, finish_time)
    
                results.append((task_id, status, exitcode, message, start_time, finish_time))

        return results

    def _write_history(self, history_db, task_id, history_id, optype):
        if not self._read_only:
            history_db.db.insert_update('fts_servers', ('url',), self.server_url)

        try:
            server_id = history_db.db.query('SELECT `id` FROM `fts_servers` WHERE `url` = %s', self.server_url)[0]
        except IndexError:
            server_id = 0

        sql = 'SELECT b.`job_id`, t.`fts_file_id` FROM `fts_{op}_tasks` AS t'
        sql += ' INNER JOIN `fts_{op}_batches` AS b ON b.`id` = t.`fts_batch_id`'
        sql += ' WHERE t.`id` = %s'

        try:
            fts_job_id, fts_file_id = self.db.query(sql.format(op = optype), task_id)[0]
        except IndexError:
            return

        if not self._read_only:
            history_db.db.insert_update('fts_batches', ('fts_server_id', 'job_id'), server_id, fts_job_id)
            batch_id = history_db.db.query('SELECT `id` FROM `fts_batches` WHERE `fts_server_id` = %s AND `job_id` = %s', server_id, fts_job_id)[0]

            history_db.db.insert_update('fts_file_{op}s'.format(op = optype), ('id', 'fts_batch_id', 'fts_file_id'), history_id, batch_id, fts_file_id)

    def _forget_status(self, task_id, optype):
        if self._read_only:
            return

        sql = 'DELETE FROM `fts_{optype}_tasks` WHERE `id` = %s'.format(optype = optype)
        self.db.query(sql, task_id)

    def _forget_batch(self, batch_id, optype):
        if self._read_only:
            return

        sql = 'DELETE FROM `fts_{optype}_batches` WHERE `batch_id` = %s'.format(optype = optype)
        self.db.query(sql, batch_id)

    def _set_server_id(self):
        if not self._read_only:
            self.db.query('INSERT INTO `fts_servers` (`url`) VALUES (%s) ON DUPLICATE KEY UPDATE `url`=VALUES(`url`)', self.server_url)

        result = self.db.query('SELECT `id` FROM `fts_servers` WHERE `url` = %s', self.server_url)
        if len(result) == 0:
            self.server_id = 0
        else:
            self.server_id = result[0]
Esempio n. 2
0
class HistoryDatabase(object):
    """
    Interface to the history database. This is a MySQL-specific implementation, and we actually
    expose the backend database.. Will be a bit tricky to replace the backend when we need to do it.
    What we do with the history DB is very much application specific, so it makes little sense
    to have generic abstract interface to individual actions. The methods of this class are just a
    few of the common operations that are necessary for any history recording.
    """

    # default configuration
    _config = Configuration()

    @staticmethod
    def set_default(config):
        HistoryDatabase._config = Configuration(config)

    def __init__(self, config=None):
        if config is None:
            config = HistoryDatabase._config

        self.db = MySQL(config.db_params)

        self.set_read_only(config.get('read_only', False))

    def set_read_only(self, value=True):
        self._read_only = value

    def save_users(self, user_list, get_ids=False):
        """
        @param user_list  [(name, dn)]
        """
        if self._read_only:
            if get_ids:
                return [0] * len(user_list)
            else:
                return

        self.db.insert_many('users', ('name', 'dn'),
                            None,
                            user_list,
                            do_update=True)

        if get_ids:
            return self.db.select_many('users', ('id', ), 'dn',
                                       [u[1] for u in user_list])

    def save_user_services(self, service_names, get_ids=False):
        if self._read_only:
            if get_ids:
                return [0] * len(service_names)
            else:
                return

        self.db.insert_many('user_services', ('name', ),
                            MySQL.make_tuple,
                            service_names,
                            do_update=True)

        if get_ids:
            return self.db.select_many('user_services', ('id', ), 'name',
                                       service_names)

    def save_partitions(self, partition_names, get_ids=False):
        if self._read_only:
            if get_ids:
                return [0] * len(partition_names)
            else:
                return

        self.db.insert_many('partitions', ('name', ),
                            MySQL.make_tuple,
                            partition_names,
                            do_update=True)

        if get_ids:
            return self.db.select_many('partitions', ('id', ), 'name',
                                       partition_names)

    def save_sites(self, site_names, get_ids=False):
        if self._read_only:
            if get_ids:
                return [0] * len(site_names)
            else:
                return

        self.db.insert_many('sites', ('name', ),
                            MySQL.make_tuple,
                            site_names,
                            do_update=True)

        if get_ids:
            return self.db.select_many('sites', ('id', ), 'name', site_names)

    def save_groups(self, group_names, get_ids=False):
        if self._read_only:
            if get_ids:
                return [0] * len(group_names)
            else:
                return

        self.db.insert_many('groups', ('name', ),
                            MySQL.make_tuple,
                            group_names,
                            do_update=True)

        if get_ids:
            return self.db.select_many('groups', ('id', ), 'name', group_names)

    def save_datasets(self, dataset_names, get_ids=False):
        if self._read_only:
            if get_ids:
                return [0] * len(dataset_names)
            else:
                return

        self.db.insert_many('datasets', ('name', ),
                            MySQL.make_tuple,
                            dataset_names,
                            do_update=True)

        if get_ids:
            return self.db.select_many('datasets', ('id', ), 'name',
                                       dataset_names)

    def save_blocks(self, block_list, get_ids=False):
        """
        @param block_list   [(dataset name, block name)]
        """
        if self._read_only:
            if get_ids:
                return [0] * len(block_list)
            else:
                return

        reuse_orig = self.db.reuse_connection
        self.db.reuse_connection = True

        datasets = set(b[0] for b in block_list)

        self.save_datasets(datasets)

        columns = [
            '`dataset` varchar(512) CHARACTER SET latin1 COLLATE latin1_general_cs NOT NULL',
            '`block` varchar(128) CHARACTER SET latin1 COLLATE latin1_general_cs NOT NULL'
        ]
        self.db.create_tmp_table('blocks_tmp', columns)
        self.db.insert_many('blocks_tmp', ('dataset', 'block'),
                            None,
                            block_list,
                            db=self.db.scratch_db)

        sql = 'INSERT INTO `blocks` (`dataset_id`, `name`)'
        sql += ' SELECT d.`id`, b.`block` FROM `{scratch}`.`blocks_tmp` AS b'.format(
            scratch=self.db.scratch_db)
        sql += ' INNER JOIN `datasets` AS d ON d.`name` = b.`dataset`'
        self.db.query(sql)

        if get_ids:
            sql = 'SELECT b.`id` FROM `blocks` AS b'
            sql += ' INNER JOIN (SELECT d.`id` dataset_id, t.`block` block_name FROM `{scratch}`.`blocks_tmp` AS t'.format(
                scratch=self.db.scratch_db)
            sql += ' INNER JOIN `datasets` AS d ON d.`name` = t.`dataset`) AS j ON (j.`dataset_id`, j.`block_name`) = (b.`dataset_id`, b.`name`)'

            ids = self.db.query(sql)

        self.db.drop_tmp_table('blocks_tmp')
        self.db.reuse_connection = reuse_orig

        if get_ids:
            return ids

    def save_files(self, file_data, get_ids=False):
        if self._read_only:
            if get_ids:
                return [0] * len(file_data)
            else:
                return

        self.db.insert_many('files', ('name', 'size'),
                            None,
                            file_data,
                            do_update=True)

        if get_ids:
            return self.db.select_many('files', ('id', ), 'name',
                                       [f[0] for f in file_data])
Esempio n. 3
0
class FTSFileOperation(FileTransferOperation, FileTransferQuery,
                       FileDeletionOperation, FileDeletionQuery):
    def __init__(self, config):
        FileTransferOperation.__init__(self, config)
        FileTransferQuery.__init__(self, config)
        FileDeletionOperation.__init__(self, config)
        FileDeletionQuery.__init__(self, config)

        self.server_url = config.fts_server
        self.server_id = 0  # server id in the DB

        # Parameter "retry" for fts3.new_job. 0 = server default
        self.fts_retry = config.get('fts_retry', 0)

        # String passed to fts3.new_*_job(metadata = _)
        self.metadata_string = config.get('metadata_string', 'Dynamo')

        # Proxy to be forwarded to FTS
        self.x509proxy = config.get('x509proxy', None)

        # Bookkeeping device
        self.db = MySQL(config.db_params)

        # Reuse the context object
        self.keep_context = config.get('keep_context', True)
        self._context = None

    def num_pending_transfers(self):  #override
        # Check the number of files in queue
        # We first thought about counting files with /files, but FTS seems to return only 1000 maximum even when "limit" is set much larger
        #files = self._ftscallurl('/files?state_in=ACTIVE,SUBMITTED,READY&limit=%d' % self.max_pending_transfers)
        #return len(files)

        num_pending = 0
        file_states = ['SUBMITTED', 'READY', 'ACTIVE', 'STAGING', 'STARTED']

        jobs = self._ftscall('list_jobs',
                             state_in=['SUBMITTED', 'ACTIVE', 'STAGING'])
        for job in jobs:
            job_info = self._ftscall('get_job_status',
                                     job['job_id'],
                                     list_files=True)
            for file_info in job_info['files']:
                if file_info['file_state'] in file_states:
                    num_pending += 1
                    if num_pending == self.max_pending_transfers + 1:
                        # don't need to query more
                        return num_pending

        return num_pending

    def num_pending_deletions(self):  #override
        # See above
        #files = self._ftscallurl('/files?state_in=ACTIVE,SUBMITTED,READY&limit=%d' % self.max_pending_deletions)
        #return len(files)

        num_pending = 0
        file_states = ['SUBMITTED', 'READY', 'ACTIVE']

        jobs = self._ftscall('list_jobs', state_in=['SUBMITTED', 'ACTIVE'])
        for job in jobs:
            job_info = self._ftscall('get_job_status',
                                     job['job_id'],
                                     list_files=True)
            for file_info in job_info['dm']:
                if file_info['file_state'] in file_states:
                    num_pending += 1
                    if num_pending == self.max_pending_deletions + 1:
                        # don't need to query more
                        return num_pending

        return num_pending

    def form_batches(self, tasks):  #override
        if len(tasks) == 0:
            return []

        # FTS3 has no restriction on how to group the transfers, but cannot apparently take thousands
        # of tasks at once
        batches = [[]]
        for task in tasks:
            batches[-1].append(task)
            if len(batches[-1]) == self.batch_size:
                batches.append([])

        return batches

    def start_transfers(self, batch_id, batch_tasks):  #override
        result = {}

        stage_files = []
        transfers = []

        s_pfn_to_task = {}
        t_pfn_to_task = {}

        for task in batch_tasks:
            sub = task.subscription
            lfn = sub.file.lfn
            dest_pfn = sub.destination.to_pfn(lfn, 'gfal2')
            source_pfn = task.source.to_pfn(lfn, 'gfal2')

            if dest_pfn is None or source_pfn is None:
                # either gfal2 is not supported or lfn could not be mapped
                LOG.warning('Could not obtain PFN for %s at %s or %s', lfn,
                            sub.destination.name, task.source.name)
                result[task] = False
                continue

            if self.checksum_algorithm:
                checksum = '%s:%s' % (
                    self.checksum_algorithm,
                    str(sub.file.checksum[self.checksum_index]))
                verify_checksum = 'target'
            else:
                checksum = None
                verify_checksum = False

            if task.source.storage_type == Site.TYPE_MSS:
                LOG.debug('Staging %s at %s', lfn, task.source.name)

                # need to stage first
                stage_files.append(
                    (source_pfn, dest_pfn, checksum, sub.file.size))

                # task identified by the source PFN
                s_pfn_to_task[source_pfn] = task
            else:
                LOG.debug('Submitting transfer of %s from %s to %s to FTS',
                          lfn, task.source.name, sub.destination.name)

                transfers.append(
                    fts3.new_transfer(source_pfn,
                                      dest_pfn,
                                      checksum=checksum,
                                      filesize=sub.file.size))

                # there should be only one task per destination pfn
                t_pfn_to_task[dest_pfn] = task

        if len(stage_files) != 0:
            LOG.debug('Submit new staging job for %d files', len(stage_files))
            job = fts3.new_staging_job([ff[0] for ff in stage_files],
                                       bring_online=36000,
                                       metadata=self.metadata_string)
            success = self._submit_job(
                job, 'staging', batch_id,
                dict(
                    (pfn, task.id) for pfn, task in s_pfn_to_task.iteritems()))

            for source_pfn, _, _, _ in stage_files:
                result[s_pfn_to_task[source_pfn]] = success

            if success and not self._read_only:
                LOG.debug('Recording staging queue')
                fields = ('id', 'source', 'destination', 'checksum', 'size')
                mapping = lambda ff: (s_pfn_to_task[ff[0]].id, ) + ff
                if not self._read_only:
                    self.db.insert_many('fts_staging_queue', fields, mapping,
                                        stage_files)

        if len(transfers) != 0:
            LOG.debug('Submit new transfer job for %d files', len(transfers))
            job = fts3.new_job(transfers,
                               retry=self.fts_retry,
                               overwrite=True,
                               verify_checksum=verify_checksum,
                               metadata=self.metadata_string)
            success = self._submit_job(
                job, 'transfer', batch_id,
                dict(
                    (pfn, task.id) for pfn, task in t_pfn_to_task.iteritems()))

            for transfer in transfers:
                dest_pfn = transfer['destinations'][0]
                result[t_pfn_to_task[dest_pfn]] = success

        return result

    def start_deletions(self, batch_id, batch_tasks):  #override
        result = {}

        pfn_to_task = {}

        for task in batch_tasks:
            desub = task.desubscription
            lfn = desub.file.lfn
            pfn = desub.site.to_pfn(lfn, 'gfal2')

            if pfn is None:
                # either gfal2 is not supported or lfn could not be mapped
                result[task] = False
                continue

            # there should be only one task per destination pfn
            pfn_to_task[pfn] = task

        job = fts3.new_delete_job(pfn_to_task.keys(),
                                  metadata=self.metadata_string)

        success = self._submit_job(
            job, 'deletion', batch_id,
            dict((pfn, task.id) for pfn, task in pfn_to_task.iteritems()))

        for task in pfn_to_task.itervalues():
            result[task] = success

        return result

    def cancel_transfers(self, task_ids):  #override
        return self._cancel(task_ids, 'transfer')

    def cancel_deletions(self, task_ids):  #override
        return self._cancel(task_ids, 'deletion')

    def cleanup(self):  #override
        sql = 'DELETE FROM f USING `fts_transfer_tasks` AS f'
        sql += ' LEFT JOIN `transfer_tasks` AS t ON t.`id` = f.`id`'
        sql += ' LEFT JOIN `fts_transfer_batches` AS b ON b.`id` = f.`fts_batch_id`'
        sql += ' WHERE t.`id` IS NULL OR b.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM f USING `fts_staging_queue` AS f'
        sql += ' LEFT JOIN `fts_transfer_tasks` AS t ON t.`id` = f.`id`'
        sql += ' WHERE t.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM f USING `fts_deletion_tasks` AS f'
        sql += ' LEFT JOIN `deletion_tasks` AS t ON t.`id` = f.`id`'
        sql += ' LEFT JOIN `fts_deletion_batches` AS b ON b.`id` = f.`fts_batch_id`'
        sql += ' WHERE t.`id` IS NULL OR b.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM f USING `fts_transfer_batches` AS f'
        sql += ' LEFT JOIN `transfer_batches` AS t ON t.`id` = f.`batch_id`'
        sql += ' WHERE t.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM f USING `fts_deletion_batches` AS f'
        sql += ' LEFT JOIN `deletion_batches` AS t ON t.`id` = f.`batch_id`'
        sql += ' WHERE t.`id` IS NULL'
        self.db.query(sql)

        # Delete the source tasks - caution: wipes out all tasks when switching the operation backend
        sql = 'DELETE FROM t USING `transfer_tasks` AS t'
        sql += ' LEFT JOIN `fts_transfer_tasks` AS f ON f.`id` = t.`id`'
        sql += ' WHERE f.`id` IS NULL'
        self.db.query(sql)
        sql = 'DELETE FROM t USING `deletion_tasks` AS t'
        sql += ' LEFT JOIN `fts_deletion_tasks` AS f ON f.`id` = t.`id`'
        sql += ' WHERE f.`id` IS NULL'
        self.db.query(sql)

    def get_transfer_status(self, batch_id):  #override
        if self.server_id == 0:
            self._set_server_id()

        results = self._get_status(batch_id, 'transfer')

        staged_tasks = []

        for task_id, status, exitcode, msg, start_time, finish_time in self._get_status(
                batch_id, 'staging'):
            if status == FileQuery.STAT_DONE:
                staged_tasks.append(task_id)
                results.append(
                    (task_id, FileQuery.STAT_QUEUED, -1, None, None, None))
            else:
                # these tasks won't appear in results from _get_status('transfer')
                # because no transfer jobs have been submitted yet
                results.append(
                    (task_id, status, exitcode, None, start_time, finish_time))

        if len(staged_tasks) != 0:
            transfers = []
            pfn_to_tid = {}
            for task_id, source_pfn, dest_pfn, checksum, filesize in self.db.select_many(
                    'fts_staging_queue',
                ('id', 'source', 'destination', 'checksum', 'size'), 'id',
                    staged_tasks):
                transfers.append(
                    fts3.new_transfer(source_pfn,
                                      dest_pfn,
                                      checksum=checksum,
                                      filesize=filesize))
                pfn_to_tid[dest_pfn] = task_id

            if self.checksum_algorithm:
                verify_checksum = 'target'
            else:
                verify_checksum = None

            job = fts3.new_job(transfers,
                               retry=self.fts_retry,
                               overwrite=True,
                               verify_checksum=verify_checksum,
                               metadata=self.metadata_string)
            success = self._submit_job(job, 'transfer', batch_id, pfn_to_tid)
            if success and not self._read_only:
                self.db.delete_many('fts_staging_queue', 'id',
                                    pfn_to_tid.values())

        return results

    def get_deletion_status(self, batch_id):  #override
        if self.server_id == 0:
            self._set_server_id()

        return self._get_status(batch_id, 'deletion')

    def write_transfer_history(self, history_db, task_id,
                               history_id):  #override
        self._write_history(history_db, task_id, history_id, 'transfer')

    def write_deletion_history(self, history_db, task_id,
                               history_id):  #override
        self._write_history(history_db, task_id, history_id, 'deletion')

    def forget_transfer_status(self, task_id):  #override
        return self._forget_status(task_id, 'transfer')

    def forget_deletion_status(self, task_id):  #override
        return self._forget_status(task_id, 'deletion')

    def forget_transfer_batch(self, task_id):  #override
        return self._forget_batch(task_id, 'transfer')

    def forget_deletion_batch(self, task_id):  #override
        return self._forget_batch(task_id, 'deletion')

    def _ftscall(self, method, *args, **kwd):
        return self._do_ftscall(binding=(method, args, kwd))

    def _ftscallurl(self, url):
        # Call to FTS URLs that don't have python bindings
        return self._do_ftscall(url=url)

    def _do_ftscall(self, binding=None, url=None):
        if self._context is None:
            # request_class = Request -> use "requests"-based https call (instead of default PyCURL,
            # which may not be able to handle proxy certificates depending on the cURL installation)
            # verify = False -> do not verify the server certificate
            context = fts3.Context(self.server_url,
                                   ucert=self.x509proxy,
                                   ukey=self.x509proxy,
                                   request_class=Request,
                                   verify=False)

            if self.keep_context:
                self._context = context
        else:
            context = self._context

        if binding is not None:
            reqstring = binding[0]
        else:
            reqstring = url

        LOG.debug('FTS: %s', reqstring)

        wait_time = 1.
        for attempt in xrange(10):
            try:
                if binding is not None:
                    method, args, kwd = binding
                    return getattr(fts3, method)(context, *args, **kwd)
                else:
                    return json.loads(context.get(url))
            except fts_exceptions.ServerError as exc:
                if str(exc.reason) == '500':
                    # Internal server error - let's try again
                    pass
            except fts_exceptions.TryAgain:
                pass

            time.sleep(wait_time)
            wait_time *= 1.5

        LOG.error('Failed to communicate with FTS server: %s', reqstring)
        raise RuntimeError('Failed to communicate with FTS server: %s' %
                           reqstring)

    def _submit_job(self, job, optype, batch_id, pfn_to_tid):
        if self._read_only:
            job_id = 'test'
        else:
            try:
                job_id = self._ftscall('submit', job)
            except:
                exc_type, exc, tb = sys.exc_info()
                LOG.error('Failed to submit %s to FTS: Exception %s (%s)',
                          optype, exc_type.__name__, str(exc))
                return False

        LOG.debug('FTS job id: %s', job_id)

        # list of file-level operations (one-to-one with pfn)
        try:
            if optype == 'transfer' or optype == 'staging':
                key = 'files'
            else:
                key = 'dm'

            fts_files = self._ftscall('get_job_status',
                                      job_id=job_id,
                                      list_files=True)[key]
        except:
            exc_type, exc, tb = sys.exc_info()
            LOG.error(
                'Failed to get status of job %s from FTS: Exception %s (%s)',
                job_id, exc_type.__name__, str(exc))
            return False

        if self.server_id == 0:
            self._set_server_id()

        if optype == 'transfer' or optype == 'staging':
            table_name = 'fts_transfer_batches'
            columns = ('batch_id', 'task_type', 'fts_server_id', 'job_id')
            values = (batch_id, optype, self.server_id, job_id)
        else:
            table_name = 'fts_deletion_batches'
            columns = ('batch_id', 'fts_server_id', 'job_id')
            values = (batch_id, self.server_id, job_id)

        if not self._read_only:
            fts_batch_id = self.db.insert_get_id(table_name,
                                                 columns=columns,
                                                 values=values)

        if optype == 'transfer' or optype == 'staging':
            table_name = 'fts_transfer_tasks'
            pfn_key = 'dest_surl'
        else:
            table_name = 'fts_deletion_tasks'
            pfn_key = 'source_surl'

        fields = ('id', 'fts_batch_id', 'fts_file_id')
        mapping = lambda f: (pfn_to_tid[f[pfn_key]], fts_batch_id, f['file_id']
                             )

        if not self._read_only:
            self.db.insert_many(table_name,
                                fields,
                                mapping,
                                fts_files,
                                do_update=True,
                                update_columns=('fts_batch_id', 'fts_file_id'))

        return True

    def _cancel(self, task_ids, optype):
        sql = 'SELECT b.`job_id`, f.`fts_file_id` FROM `fts_{op}_tasks` AS f'
        sql += ' INNER JOIN `fts_{op}_batches` AS b ON b.`id` = f.`fts_batch_id`'
        result = self.db.execute_many(sql.format(op=optype),
                                      MySQL.bare('f.`id`'), task_ids)

        by_job = collections.defaultdict(list)

        for job_id, file_id in result:
            by_job[job_id].append(file_id)

        if not self._read_only:
            for job_id, ids in by_job.iteritems():
                try:
                    self._ftscall('cancel', job_id, file_ids=ids)
                except:
                    LOG.error('Failed to cancel FTS job %s', job_id)

    def _get_status(self, batch_id, optype):
        if optype == 'transfer' or optype == 'staging':
            sql = 'SELECT `id`, `job_id` FROM `fts_transfer_batches`'
            sql += ' WHERE `task_type` = %s AND `fts_server_id` = %s AND `batch_id` = %s'
            batch_data = self.db.query(sql, optype, self.server_id, batch_id)
            task_table_name = 'fts_transfer_tasks'
        else:
            sql = 'SELECT `id`, `job_id` FROM `fts_deletion_batches`'
            sql += ' WHERE `fts_server_id` = %s AND `batch_id` = %s'
            batch_data = self.db.query(sql, self.server_id, batch_id)
            task_table_name = 'fts_deletion_tasks'

        message_pattern = re.compile(
            '(?:DESTINATION|SOURCE|TRANSFER|DELETION) \[([0-9]+)\] (.*)')

        results = []

        for fts_batch_id, job_id in batch_data:
            LOG.debug('Checking status of FTS %s batch %s', optype, job_id)

            sql = 'SELECT `fts_file_id`, `id` FROM `{table}` WHERE `fts_batch_id` = %s'.format(
                table=task_table_name)
            fts_to_task = dict(self.db.xquery(sql, fts_batch_id))

            try:
                result = self._ftscall('get_job_status',
                                       job_id=job_id,
                                       list_files=True)
            except:
                LOG.error('Failed to get job status for FTS job %s', job_id)
                continue

            if optype == 'transfer' or optype == 'staging':
                fts_files = result['files']
            else:
                fts_files = result['dm']

            for fts_file in fts_files:
                try:
                    task_id = fts_to_task[fts_file['file_id']]
                except KeyError:
                    continue

                state = fts_file['file_state']
                exitcode = -1
                start_time = None
                finish_time = None
                get_time = False

                try:
                    message = fts_file['reason']
                except KeyError:
                    message = None

                if message is not None:
                    # Check if reason follows a known format (from which we can get the exit code)
                    matches = message_pattern.match(message)
                    if matches is not None:
                        exitcode = int(matches.group(1))
                        message = matches.group(2)
                    # Additionally, if the message is a known one, convert the exit code
                    c = find_msg_code(message)
                    if c is not None:
                        exitcode = c

                    # HDFS site with gridftp-hdfs gives a I/O error (500) when the file is not there
                    if optype == 'deletion' and 'Input/output error' in message:
                        exitcode = errno.ENOENT

                if state == 'FINISHED':
                    status = FileQuery.STAT_DONE
                    exitcode = 0
                    get_time = True

                elif state == 'FAILED':
                    status = FileQuery.STAT_FAILED
                    get_time = True

                elif state == 'CANCELED':
                    status = FileQuery.STAT_CANCELLED
                    get_time = True

                elif state == 'SUBMITTED':
                    status = FileQuery.STAT_NEW

                else:
                    status = FileQuery.STAT_QUEUED

                if optype == 'transfer' and exitcode == errno.EEXIST:
                    # Transfer + destination exists -> not an error
                    status = FileQuery.STAT_DONE
                    exitcode = 0
                elif optype == 'deletion' and exitcode == errno.ENOENT:
                    # Deletion + destination does not exist -> not an error
                    status = FileQuery.STAT_DONE
                    exitcode = 0

                if get_time:
                    try:
                        start_time = calendar.timegm(
                            time.strptime(fts_file['start_time'],
                                          '%Y-%m-%dT%H:%M:%S'))
                    except TypeError:  # start time is NULL (can happen when the job is cancelled)
                        start_time = None
                    try:
                        finish_time = calendar.timegm(
                            time.strptime(fts_file['finish_time'],
                                          '%Y-%m-%dT%H:%M:%S'))
                    except TypeError:
                        start_time = None

                LOG.debug('%s %d: %s, %d, %s, %s, %s', optype, task_id,
                          FileQuery.status_name(status), exitcode, message,
                          start_time, finish_time)

                results.append((task_id, status, exitcode, message, start_time,
                                finish_time))

        return results

    def _write_history(self, history_db, task_id, history_id, optype):
        if not self._read_only:
            history_db.db.insert_update('fts_servers', ('url', ),
                                        self.server_url)

        try:
            server_id = history_db.db.query(
                'SELECT `id` FROM `fts_servers` WHERE `url` = %s',
                self.server_url)[0]
        except IndexError:
            server_id = 0

        sql = 'SELECT b.`job_id`, t.`fts_file_id` FROM `fts_{op}_tasks` AS t'
        sql += ' INNER JOIN `fts_{op}_batches` AS b ON b.`id` = t.`fts_batch_id`'
        sql += ' WHERE t.`id` = %s'

        try:
            fts_job_id, fts_file_id = self.db.query(sql.format(op=optype),
                                                    task_id)[0]
        except IndexError:
            return

        if not self._read_only:
            history_db.db.insert_update('fts_batches',
                                        ('fts_server_id', 'job_id'), server_id,
                                        fts_job_id)
            batch_id = history_db.db.query(
                'SELECT `id` FROM `fts_batches` WHERE `fts_server_id` = %s AND `job_id` = %s',
                server_id, fts_job_id)[0]

            history_db.db.insert_update('fts_file_{op}s'.format(op=optype),
                                        ('id', 'fts_batch_id', 'fts_file_id'),
                                        history_id, batch_id, fts_file_id)

    def _forget_status(self, task_id, optype):
        if self._read_only:
            return

        sql = 'DELETE FROM `fts_{optype}_tasks` WHERE `id` = %s'.format(
            optype=optype)
        self.db.query(sql, task_id)

    def _forget_batch(self, batch_id, optype):
        if self._read_only:
            return

        sql = 'DELETE FROM `fts_{optype}_batches` WHERE `batch_id` = %s'.format(
            optype=optype)
        self.db.query(sql, batch_id)

    def _set_server_id(self):
        if not self._read_only:
            self.db.query(
                'INSERT INTO `fts_servers` (`url`) VALUES (%s) ON DUPLICATE KEY UPDATE `url`=VALUES(`url`)',
                self.server_url)

        result = self.db.query(
            'SELECT `id` FROM `fts_servers` WHERE `url` = %s', self.server_url)
        if len(result) == 0:
            self.server_id = 0
        else:
            self.server_id = result[0]
Esempio n. 4
0
class MySQLReplicaLock(object):
    """
    Dataset lock read from local DB.
    Sets one attr:
      locked_blocks:   {site: set([blocks]) or None if dataset-level}
    """

    produces = ['locked_blocks']

    _default_config = None

    @staticmethod
    def set_default(config):
        MySQLReplicaLock._default_config = Configuration(config)

    def __init__(self, config = None):
        if config is None:
            if MySQLReplicaLock._default_config is None:
                raise ConfigurationError('MySQLReplicaLock default config is not set')

            config = MySQLReplicaLock._default_config

        self._mysql = MySQL(config.get('db_params', None))

        self.users = []
        for user_id, role_id in config.get('users', []):
            self.users.append((user_id, role_id))

    def load(self, inventory):
        for dataset in inventory.datasets.itervalues():
            try:
                dataset.attr.pop('locked_blocks')
            except KeyError:
                pass

        if len(self.users) != 0:
            entries = self._mysql.select_many('detox_locks', ('item', 'sites', 'groups'), ('user_id', 'role_id'), self.users)
        else:
            query = 'SELECT `item`, `sites`, `groups` FROM `detox_locks`'
            entries = self._mysql.query(query)

        for item_name, sites_pattern, groups_pattern in entries:
            # wildcard not allowed in block name
            try:
                dataset_pattern, block_name = Block.from_full_name(item_name)
            except ObjectError:
                dataset_pattern, block_name = item_name, None

            if '*' in dataset_pattern:
                pat_exp = re.compile(fnmatch.translate(dataset_pattern))
                
                datasets = []
                for dataset in inventory.datasets.values():
                    # this is highly inefficient but I can't think of a better way
                    if pat_exp.match(dataset.name):
                        datasets.append(dataset)
            else:
                try:
                    dataset = inventory.datasets[dataset_pattern]
                except KeyError:
                    LOG.debug('Cannot lock unknown dataset %s', dataset_pattern)
                    continue

                datasets = [dataset]

            specified_sites = []
            if sites_pattern:
                if sites_pattern == '*':
                    pass
                elif '*' in sites_pattern:
                    pat_exp = re.compile(fnmatch.translate(sites_pattern))
                    specified_sites.extend(s for n, s in inventory.sites.iteritems() if pat_exp.match(n))
                else:
                    try:
                        specified_sites.append(inventory.sites[sites_pattern])
                    except KeyError:
                        pass

            specified_groups = []
            if groups_pattern:
                if groups_pattern == '*':
                    pass
                elif '*' in groups_pattern:
                    pat_exp = re.compile(fnmatch.translate(groups_pattern))
                    specified_groups.extend(g for n, g in inventory.groups.iteritems() if pat_exp.match(n))
                else:
                    try:
                        specified_groups.append(inventory.groups[groups_pattern])
                    except KeyError:
                        pass

            for dataset in datasets:
                sites = set(specified_sites)
                groups = set(specified_groups)

                if len(sites) == 0:
                    # either sites_pattern was not given (global lock) or no sites matched (typo?)
                    # we will treat this as a global lock
                    sites.update(r.site for r in dataset.replicas)
    
                if len(groups) == 0:
                    # if no group matches the pattern, we will be on the safe side and treat it as a global lock
                    for replica in dataset.replicas:
                        groups.update(brep.group for brep in replica.block_replicas)
    
                try:
                    locked_blocks = dataset.attr['locked_blocks']
                except KeyError:
                    locked_blocks = dataset.attr['locked_blocks'] = {}

                if block_name is None:
                    for replica in dataset.replicas:
                        if replica.site not in sites:
                            continue
        
                        if replica.site not in locked_blocks:
                            locked_blocks[replica.site] = set()
        
                        for block_replica in replica.block_replicas:
                            if block_replica.group not in groups:
                                continue
        
                            locked_blocks[replica.site].add(block_replica.block)
                else:
                    block = dataset.find_block(block_name)
                    if block is None:
                        LOG.debug('Cannot lock unknown block %s', block_name)
                        continue

                    for replica in block.replicas:
                        if replica.site not in sites:
                            continue

                        if replica.group not in groups:
                            continue
        
                        if replica.site not in locked_blocks:
                            locked_blocks[replica.site] = set([block])
                        else:
                            locked_blocks[replica.site].add(block)
                            
        for dataset in inventory.datasets.itervalues():
            try:
                locked_blocks = dataset.attr['locked_blocks']
            except KeyError:
                continue

            for site, blocks in locked_blocks.items():
                if blocks is None:
                    continue

                # if all blocks are locked, set to None (dataset-level lock)
                if blocks == dataset.blocks:
                    locked_blocks[site] = None

        LOG.info('Locked %d items.', len(entries))
Esempio n. 5
0
class MySQLHistory(TransactionHistoryInterface):
    """
    Transaction history interface implementation using MySQL as the backend.
    """
    def __init__(self, config):
        TransactionHistoryInterface.__init__(self, config)

        self._mysql = MySQL(config.db_params)
        self._cache_db = MySQL(config.cache_db_params)

        self._site_id_map = {}
        self._dataset_id_map = {}

    def _do_acquire_lock(self, blocking):  #override
        while True:
            # Use the system table to "software-lock" the database
            self._mysql.query('LOCK TABLES `lock` WRITE')
            self._mysql.query(
                'UPDATE `lock` SET `lock_host` = %s, `lock_process` = %s WHERE `lock_host` LIKE \'\' AND `lock_process` = 0',
                socket.gethostname(), os.getpid())

            # Did the update go through?
            host, pid = self._mysql.query(
                'SELECT `lock_host`, `lock_process` FROM `lock`')[0]
            self._mysql.query('UNLOCK TABLES')

            if host == socket.gethostname() and pid == os.getpid():
                # The database is locked.
                break

            if blocking:
                LOG.warning('Failed to lock database. Waiting 30 seconds..')
                time.sleep(30)
            else:
                LOG.warning('Failed to lock database.')
                return False

        return True

    def _do_release_lock(self, force):  #override
        self._mysql.query('LOCK TABLES `lock` WRITE')
        if force:
            self._mysql.query(
                'UPDATE `lock` SET `lock_host` = \'\', `lock_process` = 0')
        else:
            self._mysql.query(
                'UPDATE `lock` SET `lock_host` = \'\', `lock_process` = 0 WHERE `lock_host` LIKE %s AND `lock_process` = %s',
                socket.gethostname(), os.getpid())

        # Did the update go through?
        host, pid = self._mysql.query(
            'SELECT `lock_host`, `lock_process` FROM `lock`')[0]
        self._mysql.query('UNLOCK TABLES')

        if host != '' or pid != 0:
            raise RuntimeError('Failed to release lock from ' +
                               socket.gethostname() + ':' + str(os.getpid()))

    def _do_new_run(self, operation, partition, policy_version,
                    comment):  #override
        part_ids = self._mysql.query(
            'SELECT `id` FROM `partitions` WHERE `name` LIKE %s', partition)
        if len(part_ids) == 0:
            part_id = self._mysql.query(
                'INSERT INTO `partitions` (`name`) VALUES (%s)', partition)
        else:
            part_id = part_ids[0]

        if operation == HistoryRecord.OP_COPY:
            if self.config.get('test', False):
                operation_str = 'copy_test'
            else:
                operation_str = 'copy'
        else:
            if self.config.get('test', False):
                operation_str = 'deletion_test'
            else:
                operation_str = 'deletion'

        return self._mysql.query(
            'INSERT INTO `runs` (`operation`, `partition_id`, `policy_version`, `comment`, `time_start`) VALUES (%s, %s, %s, %s, NOW())',
            operation_str, part_id, policy_version, comment)

    def _do_close_run(self, operation, run_number):  #override
        self._mysql.query(
            'UPDATE `runs` SET `time_end` = FROM_UNIXTIME(%s) WHERE `id` = %s',
            time.time(), run_number)

    def _do_make_copy_entry(self, run_number, site, operation_id, approved,
                            dataset_list, size):  #override
        """
        Site and datasets are expected to be already in the database.
        """

        if len(self._site_id_map) == 0:
            self._make_site_id_map()
        if len(self._dataset_id_map) == 0:
            self._make_dataset_id_map()

        self._mysql.query(
            'INSERT INTO `copy_requests` (`id`, `run_id`, `timestamp`, `approved`, `site_id`, `size`) VALUES (%s, %s, NOW(), %s, %s, %s)',
            operation_id, run_number, approved, self._site_id_map[site.name],
            size)

        self._mysql.insert_many(
            'copied_replicas', ('copy_id', 'dataset_id'), lambda d:
            (operation_id, self._dataset_id_map[d.name]), dataset_list)

    def _do_make_deletion_entry(self, run_number, site, operation_id, approved,
                                datasets, size):  #override
        """
        site and dataset are expected to be already in the database (save_deletion_decisions should be called first).
        """

        site_id = self._mysql.query(
            'SELECT `id` FROM `sites` WHERE `name` LIKE %s', site.name)[0]

        dataset_ids = self._mysql.select_many('datasets', ('id', ), 'name',
                                              (d.name for d in datasets))

        self._mysql.query(
            'INSERT INTO `deletion_requests` (`id`, `run_id`, `timestamp`, `approved`, `site_id`, `size`) VALUES (%s, %s, NOW(), %s, %s, %s)',
            operation_id, run_number, approved, site_id, size)

        self._mysql.insert_many('deleted_replicas',
                                ('deletion_id', 'dataset_id'), lambda did:
                                (operation_id, did), dataset_ids)

    def _do_update_copy_entry(self, copy_record):  #override
        self._mysql.query(
            'UPDATE `copy_requests` SET `approved` = %s, `size` = %s, `completed` = %s WHERE `id` = %s',
            copy_record.approved, copy_record.size, copy_record.completed,
            copy_record.operation_id)

    def _do_update_deletion_entry(self, deletion_record):  #override
        self._mysql.query(
            'UPDATE `deletion_requests` SET `approved` = %s, `size` = %s WHERE `id` = %s',
            deletion_record.approved, deletion_record.size,
            deletion_record.operation_id)

    def _do_save_sites(self, sites):  #override
        if len(self._site_id_map) == 0:
            self._make_site_id_map()

        names_to_sites = dict((s.name, s) for s in sites)

        sites_to_insert = []
        for site_name in names_to_sites.iterkeys():
            if site_name not in self._site_id_map:
                sites_to_insert.append(site_name)

        if len(sites_to_insert) != 0:
            self._mysql.insert_many('sites', ('name', ), None, sites_to_insert)
            self._make_site_id_map()

    def _do_get_sites(self, run_number):  #override
        self._fill_snapshot_cache('sites', run_number)

        table_name = 'sites_%d' % run_number

        sql = 'SELECT s.`name`, n.`status`, n.`quota` FROM `%s`.`%s` AS n' % (
            self._cache_db.db_name(), table_name)
        sql += ' INNER JOIN `%s`.`sites` AS s ON s.`id` = n.`site_id`' % self._mysql.db_name(
        )

        sites_dict = {}

        for site_name, status, quota in self._mysql.xquery(sql):
            sites_dict[site_name] = (status, quota)

        return sites_dict

    def _do_save_datasets(self, datasets):  #override
        if len(self._dataset_id_map) == 0:
            self._make_dataset_id_map()

        datasets_to_insert = set(d.name for d in datasets) - set(
            self._dataset_id_map.iterkeys())
        if len(datasets_to_insert) == 0:
            return

        self._mysql.insert_many('datasets', ('name', ), None,
                                datasets_to_insert)
        self._make_dataset_id_map()

    def _do_save_conditions(self, policy_lines):  #ovrride
        for line in policy_lines:
            text = re.sub('\s+', ' ', line.condition.text)
            ids = self._mysql.query(
                'SELECT `id` FROM `policy_conditions` WHERE `text` = %s', text)
            if len(ids) == 0:
                line.condition_id = self._mysql.query(
                    'INSERT INTO `policy_conditions` (`text`) VALUES (%s)',
                    text)
            else:
                line.condition_id = ids[0]

    def _do_save_copy_decisions(self, run_number, copies):  #override
        pass

    def _do_save_deletion_decisions(self, run_number, deleted_list, kept_list,
                                    protected_list):  #override
        if len(self._site_id_map) == 0:
            self._make_site_id_map()
        if len(self._dataset_id_map) == 0:
            self._make_dataset_id_map()

        if type(run_number) is int:
            # Saving deletion decisions of a cycle
            db_file_name = '%s/snapshot_%09d.db' % (
                self.config.snapshots_spool_dir, run_number)
        else:
            # run_number is actually the partition name
            db_file_name = '%s/snapshot_%s.db' % (
                self.config.snapshots_spool_dir, run_number)

        try:
            os.makedirs(self.config.snapshots_spool_dir)
            os.chmod(self.config.snapshots_spool_dir, 0777)
        except OSError:
            pass

        if os.path.exists(db_file_name):
            os.unlink(db_file_name)

        LOG.info('Creating snapshot SQLite3 DB %s', db_file_name)

        # hardcoded!!
        replica_delete = 1
        replica_keep = 2
        replica_protect = 3

        snapshot_db = sqlite3.connect(db_file_name)
        snapshot_cursor = snapshot_db.cursor()

        sql = 'CREATE TABLE `decisions` ('
        sql += '`id` TINYINT PRIMARY KEY NOT NULL,'
        sql += '`value` TEXT NOT NULL'
        sql += ')'
        snapshot_db.execute(sql)
        snapshot_db.execute('INSERT INTO `decisions` VALUES (%d, \'delete\')' %
                            replica_delete)
        snapshot_db.execute('INSERT INTO `decisions` VALUES (%d, \'keep\')' %
                            replica_keep)
        snapshot_db.execute(
            'INSERT INTO `decisions` VALUES (%d, \'protect\')' %
            replica_protect)

        sql = 'CREATE TABLE `replicas` ('
        sql += '`site_id` SMALLINT NOT NULL,'
        sql += '`dataset_id` INT NOT NULL,'
        sql += '`size` BIGINT NOT NULL,'
        sql += '`decision_id` TINYINT NOT NULL REFERENCES `decisions`(`id`),'
        sql += '`condition` MEDIUMINT NOT NULL'
        sql += ')'
        snapshot_db.execute(sql)
        snapshot_db.execute(
            'CREATE INDEX `site_dataset` ON `replicas` (`site_id`, `dataset_id`)'
        )

        sql = 'INSERT INTO `replicas` VALUES (?, ?, ?, ?, ?)'

        def do_insert(entries, decision):
            for replica, matches in entries.iteritems():
                site_id = self._site_id_map[replica.site.name]
                dataset_id = self._dataset_id_map[replica.dataset.name]

                for condition_id, block_replicas in matches.iteritems():
                    size = sum(r.size for r in block_replicas)
                    snapshot_cursor.execute(
                        sql,
                        (site_id, dataset_id, size, decision, condition_id))

            snapshot_db.commit()

        do_insert(deleted_list, replica_delete)
        do_insert(kept_list, replica_keep)
        do_insert(protected_list, replica_protect)

        snapshot_cursor.close()
        snapshot_db.close()

        os.chmod(db_file_name, 0666)

        self._fill_snapshot_cache('replicas', run_number, overwrite=True)

    def _do_save_quotas(self, run_number, quotas):  #override
        # Will save quotas and statuses

        if len(self._site_id_map) == 0:
            self._make_site_id_map()

        if type(run_number) is int:
            db_file_name = '%s/snapshot_%09d.db' % (
                self.config.snapshots_spool_dir, run_number)
        else:
            # run_number is actually the partition name
            db_file_name = '%s/snapshot_%s.db' % (
                self.config.snapshots_spool_dir, run_number)

        # DB file should exist already - this function is called after save_deletion_decisions

        snapshot_db = sqlite3.connect(db_file_name)
        snapshot_cursor = snapshot_db.cursor()

        sql = 'CREATE TABLE `statuses` ('
        sql += '`id` TINYINT PRIMARY KEY NOT NULL,'
        sql += '`value` TEXT NOT NULL'
        sql += ')'
        snapshot_db.execute(sql)
        snapshot_db.execute('INSERT INTO `statuses` VALUES (%d, \'ready\')' %
                            Site.STAT_READY)
        snapshot_db.execute(
            'INSERT INTO `statuses` VALUES (%d, \'waitroom\')' %
            Site.STAT_WAITROOM)
        snapshot_db.execute('INSERT INTO `statuses` VALUES (%d, \'morgue\')' %
                            Site.STAT_MORGUE)
        snapshot_db.execute('INSERT INTO `statuses` VALUES (%d, \'unknown\')' %
                            Site.STAT_UNKNOWN)

        sql = 'CREATE TABLE `sites` ('
        sql += '`site_id` SMALLINT PRIMARY KEY NOT NULL,'
        sql += '`status_id` TINYINT NOT NULL REFERENCES `statuses`(`id`),'
        sql += '`quota` INT NOT NULL'
        sql += ')'
        snapshot_db.execute(sql)

        sql = 'INSERT INTO `sites` VALUES (?, ?, ?)'

        for site, quota in quotas.iteritems():
            snapshot_cursor.execute(
                sql, (self._site_id_map[site.name], site.status, quota))

        snapshot_db.commit()

        snapshot_cursor.close()
        snapshot_db.close()

        self._fill_snapshot_cache('sites', run_number, overwrite=True)

        if type(run_number) is int:
            # This was a numbered cycle
            # Archive the sqlite3 file
            # Relying on the fact save_quotas is called after save_deletion_decisions

            srun = '%09d' % run_number
            archive_dir_name = '%s/%s/%s' % (self.config.snapshots_archive_dir,
                                             srun[:3], srun[3:6])
            xz_file_name = '%s/snapshot_%09d.db.xz' % (archive_dir_name,
                                                       run_number)

            try:
                os.makedirs(archive_dir_name)
            except OSError:
                pass

            with open(db_file_name, 'rb') as db_file:
                with open(xz_file_name, 'wb') as xz_file:
                    xz_file.write(lzma.compress(db_file.read()))

    def _do_get_deletion_decisions(self, run_number, size_only):  #override
        self._fill_snapshot_cache('replicas', run_number)

        table_name = 'replicas_%d' % run_number

        if size_only:
            # return {site_name: (protect_size, delete_size, keep_size)}
            volumes = {}
            sites = set()

            query = 'SELECT s.`name`, SUM(r.`size`) * 1.e-12 FROM `%s`.`%s` AS r' % (
                self._cache_db.db_name(), table_name)
            query += ' INNER JOIN `%s`.`sites` AS s ON s.`id` = r.`site_id`' % self._mysql.db_name(
            )
            query += ' WHERE r.`decision` LIKE %s'
            query += ' GROUP BY r.`site_id`'

            for decision in ['protect', 'delete', 'keep']:
                volumes[decision] = dict(self._mysql.xquery(query, decision))
                sites.update(set(volumes[decision].iterkeys()))

            product = {}
            for site_name in sites:
                v = {}
                for decision in ['protect', 'delete', 'keep']:
                    try:
                        v[decision] = volumes[decision][site_name]
                    except:
                        v[decision] = 0

                product[site_name] = (v['protect'], v['delete'], v['keep'])

            return product

        else:
            # return {site_name: [(dataset_name, size, decision, reason)]}

            query = 'SELECT s.`name`, d.`name`, r.`size`, r.`decision`, p.`text` FROM `%s`.`%s` AS r' % (
                self._cache_db.db_name(), table_name)
            query += ' INNER JOIN `%s`.`sites` AS s ON s.`id` = r.`site_id`' % self._mysql.db_name(
            )
            query += ' INNER JOIN `%s`.`datasets` AS d ON d.`id` = r.`dataset_id`' % self._mysql.db_name(
            )
            query += ' INNER JOIN `%s`.`policy_conditions` AS p ON p.`id` = r.`condition`' % self._mysql.db_name(
            )
            query += ' ORDER BY s.`name` ASC, r.`size` DESC'

            product = {}

            _site_name = ''

            for site_name, dataset_name, size, decision, reason in self._cache_db.xquery(
                    query):
                if site_name != _site_name:
                    product[site_name] = []
                    current = product[site_name]
                    _site_name = site_name

                current.append((dataset_name, size, decision, reason))

            return product

    def _do_save_dataset_popularity(self, run_number, datasets):  #override
        if len(self._dataset_id_map) == 0:
            self._make_dataset_id_map()

        fields = ('run_id', 'dataset_id', 'popularity')
        mapping = lambda dataset: (run_number, self._dataset_id_map[
            dataset.name], dataset.attr['request_weight']
                                   if 'request_weight' in dataset.attr else 0.)
        self._mysql.insert_many('dataset_popularity_snapshots', fields,
                                mapping, datasets)

    def _do_get_incomplete_copies(self, partition):  #override
        query = 'SELECT h.`id`, UNIX_TIMESTAMP(h.`timestamp`), h.`approved`, s.`name`, h.`size`'
        query += ' FROM `copy_requests` AS h'
        query += ' INNER JOIN `runs` AS r ON r.`id` = h.`run_id`'
        query += ' INNER JOIN `partitions` AS p ON p.`id` = r.`partition_id`'
        query += ' INNER JOIN `sites` AS s ON s.`id` = h.`site_id`'
        query += ' WHERE h.`id` > 0 AND p.`name` LIKE \'%s\' AND h.`completed` = 0 AND h.`run_id` > 0' % partition
        history_entries = self._mysql.xquery(query)

        id_to_record = {}
        for eid, timestamp, approved, site_name, size in history_entries:
            id_to_record[eid] = HistoryRecord(HistoryRecord.OP_COPY,
                                              eid,
                                              site_name,
                                              timestamp=timestamp,
                                              approved=approved,
                                              size=size)

        id_to_dataset = dict(
            self._mysql.xquery('SELECT `id`, `name` FROM `datasets`'))
        id_to_site = dict(
            self._mysql.xquery('SELECT `id`, `name` FROM `sites`'))

        replicas = self._mysql.select_many('copied_replicas',
                                           ('copy_id', 'dataset_id'),
                                           'copy_id', id_to_record.iterkeys())

        current_copy_id = 0
        for copy_id, dataset_id in replicas:
            if copy_id != current_copy_id:
                record = id_to_record[copy_id]
                current_copy_id = copy_id

            record.replicas.append(
                HistoryRecord.CopiedReplica(
                    dataset_name=id_to_dataset[dataset_id]))

        return id_to_record.values()

    def _do_get_copied_replicas(self, run_number):  #override
        query = 'SELECT s.`name`, d.`name` FROM `copied_replicas` AS p'
        query += ' INNER JOIN `copy_requests` AS r ON r.`id` = p.`copy_id`'
        query += ' INNER JOIN `datasets` AS d ON d.`id` = p.`dataset_id`'
        query += ' INNER JOIN `sites` AS s ON s.`id` = r.`site_id`'
        query += ' WHERE r.`run_id` = %d' % run_number

        return self._mysql.query(query)

    def _do_get_site_name(self, operation_id):  #override
        result = self._mysql.query(
            'SELECT s.name FROM `sites` AS s INNER JOIN `copy_requests` AS h ON h.`site_id` = s.`id` WHERE h.`id` = %s',
            operation_id)
        if len(result) != 0:
            return result[0]

        result = self._mysql.query(
            'SELECT s.name FROM `sites` AS s INNER JOIN `deletion_requests` AS h ON h.`site_id` = s.`id` WHERE h.`id` = %s',
            operation_id)
        if len(result) != 0:
            return result[0]

        return ''

    def _do_get_deletion_runs(self, partition, first, last):  #override
        result = self._mysql.query(
            'SELECT `id` FROM `partitions` WHERE `name` LIKE %s', partition)
        if len(result) == 0:
            return []

        partition_id = result[0]

        sql = 'SELECT `id` FROM `runs` WHERE `partition_id` = %d AND `time_end` NOT LIKE \'0000-00-00 00:00:00\' AND `operation` IN (\'deletion\', \'deletion_test\')' % partition_id

        if first >= 0:
            sql += ' AND `id` >= %d' % first
        if last >= 0:
            sql += ' AND `id` <= %d' % last

        sql += ' ORDER BY `id` ASC'

        result = self._mysql.query(sql)

        if first < 0 and len(result) > 1:
            result = result[-1:]

        return result

    def _do_get_copy_runs(self, partition, first, last):  #override
        result = self._mysql.query(
            'SELECT `id` FROM `partitions` WHERE `name` LIKE %s', partition)
        if len(result) == 0:
            return []

        partition_id = result[0]

        sql = 'SELECT `id` FROM `runs` WHERE `partition_id` = %d AND `time_end` NOT LIKE \'0000-00-00 00:00:00\' AND `operation` IN (\'copy\', \'copy_test\')' % partition_id

        if first >= 0:
            sql += ' AND `id` >= %d' % first
        if last >= 0:
            sql += ' AND `id` <= %d' % last

        sql += ' ORDER BY `id` ASC'

        if first < 0 and len(result) > 1:
            result = result[-1:]

        return result

    def _do_get_run_timestamp(self, run_number):  #override
        result = self._mysql.query(
            'SELECT UNIX_TIMESTAMP(`time_start`) FROM `runs` WHERE `id` = %s',
            run_number)
        if len(result) == 0:
            return 0

        return result[0]

    def _do_get_next_test_id(self):  #override
        copy_result = self._mysql.query(
            'SELECT MIN(`id`) FROM `copy_requests`')[0]
        if copy_result == None:
            copy_result = 0

        deletion_result = self._mysql.query(
            'SELECT MIN(`id`) FROM `deletion_requests`')[0]
        if deletion_result == None:
            deletion_result = 0

        return min(copy_result, deletion_result) - 1

    def _make_site_id_map(self):
        self._site_id_map = {}
        for name, site_id in self._mysql.xquery(
                'SELECT `name`, `id` FROM `sites`'):
            self._site_id_map[name] = int(site_id)

    def _make_dataset_id_map(self):
        self._dataset_id_map = {}
        for name, dataset_id in self._mysql.xquery(
                'SELECT `name`, `id` FROM `datasets`'):
            self._dataset_id_map[name] = int(dataset_id)

    def _fill_snapshot_cache(self, template, run_number, overwrite=False):
        # run_number is either a cycle number or a partition name. %s works for both
        table_name = '%s_%s' % (template, run_number)

        sql = 'SELECT COUNT(*) FROM `information_schema`.`TABLES` WHERE `TABLE_SCHEMA` = %s AND `TABLE_NAME` = %s'
        table_exists = (self._mysql.query(sql, self._cache_db.db_name(),
                                          table_name)[0] != 0)

        if overwrite or not table_exists:
            # fill from sqlite
            if table_exists:
                self._cache_db.query('TRUNCATE TABLE `%s`' % table_name)
            else:
                self._cache_db.query('CREATE TABLE `%s` LIKE `%s`' %
                                     (table_name, template))

            if type(run_number) is int:
                db_file_name = '%s/snapshot_%09d.db' % (
                    self.config.snapshots_spool_dir, run_number)

                if not os.path.exists(db_file_name):
                    srun = '%09d' % run_number
                    xz_file_name = '%s/%s/%s/snapshot_%09d.db.xz' % (
                        self.config.snapshots_archive_dir, srun[:3], srun[3:6],
                        run_number)
                    if not os.path.exists(xz_file_name):
                        raise RuntimeError('Snapshot DB ' + db_file_name +
                                           ' does not exist')

                    with open(xz_file_name, 'rb') as xz_file:
                        with open(db_file_name, 'wb') as db_file:
                            db_file.write(lzma.decompress(xz_file.read()))

            else:
                db_file_name = '%s/snapshot_%s.db' % (
                    self.config.snapshots_spool_dir, run_number)

                if not os.path.exists(db_file_name):
                    return

            snapshot_db = sqlite3.connect(db_file_name)
            snapshot_db.text_factory = str  # otherwise we'll get unicode and MySQLdb cannot convert that
            snapshot_cursor = snapshot_db.cursor()

            def make_snapshot_reader():
                if template == 'replicas':
                    sql = 'SELECT r.`site_id`, r.`dataset_id`, r.`size`, d.`value`, r.`condition` FROM `replicas` AS r'
                    sql += ' INNER JOIN `decisions` AS d ON d.`id` = r.`decision_id`'
                elif template == 'sites':
                    sql = 'SELECT s.`site_id`, t.`value`, s.`quota` FROM `sites` AS s'
                    sql += ' INNER JOIN `statuses` AS t ON t.`id` = s.`status_id`'

                snapshot_cursor.execute(sql)

                while True:
                    row = snapshot_cursor.fetchone()
                    if row is None:
                        return

                    yield row

            snapshot_reader = make_snapshot_reader()

            if template == 'replicas':
                fields = ('site_id', 'dataset_id', 'size', 'decision',
                          'condition')
            elif template == 'sites':
                fields = ('site_id', 'status', 'quota')

            self._cache_db.insert_many(table_name,
                                       fields,
                                       None,
                                       snapshot_reader,
                                       do_update=False)

            snapshot_cursor.close()
            snapshot_db.close()

        if type(run_number) is int:
            self._cache_db.query(
                'INSERT INTO `{template}_snapshot_usage` VALUES (%s, NOW())'.
                format(template=template), run_number)

            # also save into the main cache table
            sql = 'SELECT p.`name` FROM `partitions` AS p INNER JOIN `runs` AS r ON r.`partition_id` = p.`id` WHERE r.`id` = %s'
            partition = self._mysql.query(sql, run_number)[0]

            self._fill_snapshot_cache(template, partition, overwrite)

        self._clean_old_cache()

    def _clean_old_cache(self):
        sql = 'SELECT `run_id` FROM (SELECT `run_id`, MAX(`timestamp`) AS m FROM `replicas_snapshot_usage` GROUP BY `run_id`) AS t WHERE m < DATE_SUB(NOW(), INTERVAL 1 WEEK)'
        old_replica_runs = self._cache_db.query(sql)
        for old_run in old_replica_runs:
            table_name = 'replicas_%d' % old_run
            self._cache_db.query('DROP TABLE IF EXISTS `%s`' % table_name)

        sql = 'SELECT `run_id` FROM (SELECT `run_id`, MAX(`timestamp`) AS m FROM `sites_snapshot_usage` GROUP BY `run_id`) AS t WHERE m < DATE_SUB(NOW(), INTERVAL 1 WEEK)'
        old_site_runs = self._cache_db.query(sql)
        for old_run in old_site_runs:
            table_name = 'sites_%d' % old_run
            self._cache_db.query('DROP TABLE IF EXISTS `%s`' % table_name)

        for old_run in set(old_replica_runs) & set(old_site_runs):
            srun = '%09d' % old_run
            db_file_name = '%s/snapshot_%09d.db' % (
                self.config.snapshots_spool_dir, old_run)
            if os.path.exists(db_file_name):
                try:
                    os.unlink(db_file_name)
                except:
                    LOG.error('Failed to delete %s' % db_file_name)
                    pass

        self._cache_db.query(
            'DELETE FROM `replicas_snapshot_usage` WHERE `timestamp` < DATE_SUB(NOW(), INTERVAL 1 WEEK)'
        )
        self._cache_db.query('OPTIMIZE TABLE `replicas_snapshot_usage`')
        self._cache_db.query(
            'DELETE FROM `sites_snapshot_usage` WHERE `timestamp` < DATE_SUB(NOW(), INTERVAL 1 WEEK)'
        )
        self._cache_db.query('OPTIMIZE TABLE `sites_snapshot_usage`')
Esempio n. 6
0
class HistoryDatabase(object):
    """
    Interface to the history database. This is a MySQL-specific implementation, and we actually
    expose the backend database.. Will be a bit tricky to replace the backend when we need to do it.
    What we do with the history DB is very much application specific, so it makes little sense
    to have generic abstract interface to individual actions. The methods of this class are just a
    few of the common operations that are necessary for any history recording.
    """

    # default configuration
    _config = Configuration()

    @staticmethod
    def set_default(config):
        HistoryDatabase._config = Configuration(config)

    def __init__(self, config = None):
        if config is None:
            config = HistoryDatabase._config

        self.db = MySQL(config.db_params)

        self.set_read_only(config.get('read_only', False))

    def set_read_only(self, value = True):
        self._read_only = value

    def save_users(self, user_list, get_ids = False):
        """
        @param user_list  [(name, dn)]
        """
        if self._read_only:
            if get_ids:
                return [0] * len(user_list)
            else:
                return

        self.db.insert_many('users', ('name', 'dn'), None, user_list, do_update = True)

        if get_ids:
            return self.db.select_many('users', ('id',), 'dn', [u[1] for u in user_list])

    def save_user_services(self, service_names, get_ids = False):
        if self._read_only:
            if get_ids:
                return [0] * len(service_names)
            else:
                return

        self.db.insert_many('user_services', ('name',), MySQL.make_tuple, service_names, do_update = True)

        if get_ids:
            return self.db.select_many('user_services', ('id',), 'name', service_names)

    def save_partitions(self, partition_names, get_ids = False):
        if self._read_only:
            if get_ids:
                return [0] * len(partition_names)
            else:
                return

        self.db.insert_many('partitions', ('name',), MySQL.make_tuple, partition_names, do_update = True)

        if get_ids:
            return self.db.select_many('partitions', ('id',), 'name', partition_names)

    def save_sites(self, site_names, get_ids = False):
        if self._read_only:
            if get_ids:
                return [0] * len(site_names)
            else:
                return

        self.db.insert_many('sites', ('name',), MySQL.make_tuple, site_names, do_update = True)

        if get_ids:
            return self.db.select_many('sites', ('id',), 'name', site_names)

    def save_groups(self, group_names, get_ids = False):
        if self._read_only:
            if get_ids:
                return [0] * len(group_names)
            else:
                return

        self.db.insert_many('groups', ('name',), MySQL.make_tuple, group_names, do_update = True)

        if get_ids:
            return self.db.select_many('groups', ('id',), 'name', group_names)

    def save_datasets(self, dataset_names, get_ids = False):
        if self._read_only:
            if get_ids:
                return [0] * len(dataset_names)
            else:
                return

        self.db.insert_many('datasets', ('name',), MySQL.make_tuple, dataset_names, do_update = True)

        if get_ids:
            return self.db.select_many('datasets', ('id',), 'name', dataset_names)

    def save_blocks(self, block_list, get_ids = False):
        """
        @param block_list   [(dataset name, block name)]
        """
        if self._read_only:
            if get_ids:
                return [0] * len(block_list)
            else:
                return

        reuse_orig = self.db.reuse_connection
        self.db.reuse_connection = True

        datasets = set(b[0] for b in block_list)

        self.save_datasets(datasets)

        columns = [
            '`dataset` varchar(512) CHARACTER SET latin1 COLLATE latin1_general_cs NOT NULL',
            '`block` varchar(128) CHARACTER SET latin1 COLLATE latin1_general_cs NOT NULL'
        ]
        self.db.create_tmp_table('blocks_tmp', columns)
        self.db.insert_many('blocks_tmp', ('dataset', 'block'), None, block_list, db = self.db.scratch_db)

        sql = 'INSERT INTO `blocks` (`dataset_id`, `name`)'
        sql += ' SELECT d.`id`, b.`block` FROM `{scratch}`.`blocks_tmp` AS b'.format(scratch = self.db.scratch_db)
        sql += ' INNER JOIN `datasets` AS d ON d.`name` = b.`dataset`'
        self.db.query(sql)

        if get_ids:
            sql = 'SELECT b.`id` FROM `blocks` AS b'
            sql += ' INNER JOIN (SELECT d.`id` dataset_id, t.`block` block_name FROM `{scratch}`.`blocks_tmp` AS t'.format(scratch = self.db.scratch_db)
            sql += ' INNER JOIN `datasets` AS d ON d.`name` = t.`dataset`) AS j ON (j.`dataset_id`, j.`block_name`) = (b.`dataset_id`, b.`name`)'

            ids = self.db.query(sql)

        self.db.drop_tmp_table('blocks_tmp')
        self.db.reuse_connection = reuse_orig

        if get_ids:
            return ids

    def save_files(self, file_data, get_ids = False):
        if self._read_only:
            if get_ids:
                return [0] * len(file_data)
            else:
                return

        self.db.insert_many('files', ('name', 'size'), None, file_data, do_update = True)

        if get_ids:
            return self.db.select_many('files', ('id',), 'name', [f[0] for f in file_data])