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
0015 args=parser.parse_args()
0016
0017 file_paths = args.file_paths
0018 did_names = args.did_names
0019 scope= args.scope
0020 rse= args.rse
0021
0022 # Validation to ensure file_paths and did_names have the same length
0023 if len(file_paths) != len(did_names):
0024     raise ValueError("The number of file paths must match the number of did names.")
0025
0026 upload_items = []  # List to hold the upload items
0027
0028 # Loop through the file paths and did names (assuming did_names length matches file_paths length)
0029 for file_path, did_name in zip(file_paths, did_names):
0030     parent_directory = os.path.dirname(did_name)  # Get the parent directory from did_name
0031
0032     # Create a new dictionary for each file and did_name
0033     upload_item = {
0034         'path': file_path,
0035         'rse': rse,
0036         'did_scope': scope,
0037         'did_name': did_name,
0038         'dataset_scope': scope,
0039         'dataset_name': parent_directory
0040     }
0041
0042     # Append the new item to the upload_items list
0043     upload_items.append(upload_item)
0044
0045 logger = logging.getLogger('upload_client')
0046 logger.addHandler(logging.StreamHandler())
0047 logger.setLevel(logging.INFO)
0048 upload_client=UploadClient(logger=logger)
0049 upload_client.upload(upload_items)