simulation_campaign_hepmc3/scripts/register_to_rucio.py

0001 #!/usr/bin/env python3
0002
0003 import argparse
0004 import os
0005 from rucio.client.uploadclient import UploadClient
0006 from rucio.common.exception import InputValidationError, RSEWriteBlocked, NoFilesUploaded, NotAllFilesUploaded
0007 import logging
0008
0009 parser = argparse.ArgumentParser(prog='Register to RUCIO', description='Registers files to RUCIO')
0010 parser.add_argument("-f", dest="file_paths", action="store", nargs='+', required=True, help="Enter the local file path")
0011 parser.add_argument("-d", dest="did_names", action="store", nargs='+', required=True, help="Enter the data identifier for rucio catalogue")
0012 parser.add_argument("-s", dest="scope", action="store", required=True, help="Enter the scope")
0013 parser.add_argument("-r", dest="rse", action="store", required=True, help="Enter the rucio storage element. EIC-XRD is for storing production outputs.")
0014 parser.add_argument('-noregister', dest="noregister", action="store_true", default=False, help="Specify if rucio registration should be skipped")
0015
0016 args=parser.parse_args()
0017
0018 file_paths = args.file_paths
0019 did_names = args.did_names
0020 scope= args.scope
0021 rse= args.rse
0022 noregister = args.noregister
0023
0024 # Validation to ensure file_paths and did_names have the same length
0025 if len(file_paths) != len(did_names):
0026     raise ValueError("The number of file paths must match the number of did names.")
0027
0028 upload_items = []  # List to hold the upload items
0029
0030 # Loop through the file paths and did names (assuming did_names length matches file_paths length)
0031 for file_path, did_name in zip(file_paths, did_names):
0032     parent_directory = os.path.dirname(did_name)  # Get the parent directory from did_name
0033
0034     # Create a new dictionary for each file and did_name
0035     upload_item = {
0036         'path': file_path,
0037         'rse': rse,
0038         'did_scope': scope,
0039         'did_name': did_name,
0040         'dataset_scope': scope,
0041         'dataset_name': parent_directory,
0042         'no_register': noregister
0043     }
0044
0045     # Append the new item to the upload_items list
0046     upload_items.append(upload_item)
0047
0048 logger = logging.getLogger('upload_client')
0049 logger.addHandler(logging.StreamHandler())
0050 logger.setLevel(logging.INFO)
0051 upload_client=UploadClient(logger=logger)
0052 upload_client.upload(upload_items)