|
| 1 | +# pylint: disable=consider-using-f-string, unspecified-encoding |
| 2 | +""" |
| 3 | +allow users to (in)validate some files in their USER datasets in phys03 |
| 4 | +""" |
| 5 | + |
| 6 | +import json |
| 7 | + |
| 8 | +from CRABClient.Commands.SubCommand import SubCommand |
| 9 | +from CRABClient.ClientExceptions import MissingOptionException, ConfigurationException, CommandFailedException |
| 10 | +from CRABClient.ClientUtilities import colors |
| 11 | +from CRABClient.RestInterfaces import getDbsREST |
| 12 | + |
| 13 | + |
| 14 | +class setfilestatus(SubCommand): |
| 15 | + """ |
| 16 | + Set status of a USER dataset in phys03, |
| 17 | + optionally invalidates/revalidates all files in it |
| 18 | + meant to replace https://github.com/dmwm/DBS/blob/master/Client/utils/DataOpsScripts/DBS3SetDatasetStatus.py |
| 19 | + and to work whenever CRAB is supported, i.e. with both python2 and python3 |
| 20 | + """ |
| 21 | + |
| 22 | + name = 'setfilestatus' |
| 23 | + |
| 24 | + def __init__(self, logger, cmdargs=None): |
| 25 | + SubCommand.__init__(self, logger, cmdargs) |
| 26 | + |
| 27 | + def __call__(self): |
| 28 | + |
| 29 | + result = 'FAILED' # will change to 'SUCCESS' when all is OK |
| 30 | + |
| 31 | + # intitalize, and validate args |
| 32 | + instance = self.options.instance |
| 33 | + dataset = self.options.dataset |
| 34 | + files = self.options.files |
| 35 | + status = self.options.status |
| 36 | + self.logger.debug('instance = %s' % instance) |
| 37 | + self.logger.debug('dataset = %s' % dataset) |
| 38 | + self.logger.debug('files = %s' % files) |
| 39 | + self.logger.debug('status = %s' % status) |
| 40 | + |
| 41 | + statusToSet = 1 if status == 'VALID' else 0 |
| 42 | + |
| 43 | + filesToChange = None |
| 44 | + if files: |
| 45 | + # did the user specify the name of a file containing a list of LFN's ? |
| 46 | + try: |
| 47 | + with open(files, 'r') as f: |
| 48 | + flist = [lfn.strip() for lfn in f] |
| 49 | + filesToChange = ','.join(flist) |
| 50 | + except IOError: |
| 51 | + # no. Assume we have a comma separated list of LFN's (a single LFN is also OK) |
| 52 | + filesToChange = files.strip(",").strip() |
| 53 | + finally: |
| 54 | + # files and dataset options are mutually exclusive |
| 55 | + dataset = None |
| 56 | + if ',' in filesToChange: |
| 57 | + raise NotImplementedError('list of LFNs is not supported yet') |
| 58 | + |
| 59 | + # from DBS instance, to DBS REST services |
| 60 | + dbsReader, dbsWriter = getDbsREST(instance=instance, logger=self.logger, |
| 61 | + cert=self.proxyfilename, key=self.proxyfilename) |
| 62 | + # we will need the dataset name |
| 63 | + if dataset: |
| 64 | + datasetName = dataset |
| 65 | + else: |
| 66 | + # get it from DBS |
| 67 | + lfn = filesToChange.split(',')[0] |
| 68 | + query = {'logical_file_name': lfn} |
| 69 | + out, rc, msg = dbsReader.get(uri='datasets', data=query) |
| 70 | + if not out: |
| 71 | + self.logger.error("ERROR: file %s not found in DBS" % lfn) |
| 72 | + raise ConfigurationException |
| 73 | + datasetName = out[0]['dataset'] |
| 74 | + self.logger.info('LFN to be changed belongs to dataset %s' % datasetName) |
| 75 | + |
| 76 | + # when acting on a list of LFN's, can't print status of all files before/after |
| 77 | + # best we can do is to print the number of valid/invalid file in the dataset |
| 78 | + # before/after. |
| 79 | + |
| 80 | + self.logFilesTally(dataset=datasetName, dbs=dbsReader) |
| 81 | + |
| 82 | + if filesToChange: |
| 83 | + data = {'logical_file_name': filesToChange, 'is_file_valid': statusToSet} |
| 84 | + if dataset: |
| 85 | + data = {'dataset': dataset, 'is_file_valid': statusToSet} |
| 86 | + jdata = json.dumps(data) # PUT requires data in JSON format |
| 87 | + out, rc, msg = dbsWriter.put(uri='files', data=jdata) |
| 88 | + if rc == 200 and msg == 'OK': |
| 89 | + self.logger.info("File(s) status changed successfully") |
| 90 | + result = 'SUCCESS' |
| 91 | + else: |
| 92 | + msg = "File(s) status change failed: %s" % out |
| 93 | + raise CommandFailedException(msg) |
| 94 | + |
| 95 | + self.logFilesTally(dataset=datasetName, dbs=dbsReader) |
| 96 | + |
| 97 | + return {'commandStatus': result} |
| 98 | + |
| 99 | + def logFilesTally(self, dataset=None, dbs=None): |
| 100 | + """ prints total/valid/invalid files in dataset """ |
| 101 | + query = {'dataset': dataset, 'validFileOnly': 1} |
| 102 | + out, _, _ = dbs.get(uri='files', data=query) |
| 103 | + valid = len(out) |
| 104 | + query = {'dataset': dataset, 'validFileOnly': 0} |
| 105 | + out, _, _ = dbs.get(uri='files', data=query) |
| 106 | + total = len(out) |
| 107 | + invalid = total - valid |
| 108 | + self.logger.info("Dataset file count total/valid/invalid = %d/%d/%d" % (total, valid, invalid)) |
| 109 | + |
| 110 | + def setOptions(self): |
| 111 | + """ |
| 112 | + __setOptions__ |
| 113 | +
|
| 114 | + This allows to set specific command options |
| 115 | + """ |
| 116 | + self.parser.add_option('-i', '--instance', dest='instance', default='prod/phys03', |
| 117 | + help='DBS instance. e.g. prod/phys03 (default) or int/phys03' |
| 118 | + ) |
| 119 | + self.parser.add_option('-d', '--dataset', dest='dataset', default=None, |
| 120 | + help='Will apply status to all files in this dataset.' + \ |
| 121 | + ' Use either --files or--dataset', |
| 122 | + metavar='<dataset_name>') |
| 123 | + self.parser.add_option('-s', '--status', dest='status', default=None, |
| 124 | + help='New status of the file(s): VALID/INVALID', |
| 125 | + choices=['VALID', 'INVALID'] |
| 126 | + ) |
| 127 | + self.parser.add_option('-f', '--files', dest='files', default=None, |
| 128 | + help='List of files to be validated/invalidated.' + \ |
| 129 | + ' Can be either a simple LFN or a file containg LFNs or' + \ |
| 130 | + ' a comma separated list of LFNs. Use either --files or --dataset', |
| 131 | + metavar="<lfn1[,..,lfnx] or filename>") |
| 132 | + |
| 133 | + def validateOptions(self): |
| 134 | + SubCommand.validateOptions(self) |
| 135 | + |
| 136 | + if not self.options.files and not self.options.dataset: |
| 137 | + msg = "%sError%s: Please specify the files to change." % (colors.RED, colors.NORMAL) |
| 138 | + msg += " Use either the --files or the --dataset option." |
| 139 | + ex = MissingOptionException(msg) |
| 140 | + ex.missingOption = "files" |
| 141 | + raise ex |
| 142 | + if self.options.files and self.options.dataset: |
| 143 | + msg = "%sError%s: You can not use both --files and --dataset at same time" % (colors.RED, colors.NORMAL) |
| 144 | + raise ConfigurationException(msg) |
| 145 | + if self.options.status is None: |
| 146 | + msg = "%sError%s: Please specify the new file(s) status." % (colors.RED, colors.NORMAL) |
| 147 | + msg += " Use the --status option." |
| 148 | + ex = MissingOptionException(msg) |
| 149 | + ex.missingOption = "status" |
| 150 | + raise ex |
0 commit comments