src/swf_common_lib/rucio_utils.py

0001 """
0002 Utility functions for Rucio workflow operations.
0003 """
0004
0005 import hashlib
0006 import logging
0007 import os
0008 import re
0009 import zlib
0010 from typing import Optional
0011
0012 from rucio.client import Client as RucioClient
0013 from rucio.common.exception import (
0014     DataIdentifierAlreadyExists,
0015     DataIdentifierNotFound,
0016     FileAlreadyExists,
0017     RSENotFound,
0018 )
0019
0020
0021 logger = logging.getLogger(__name__)
0022
0023
0024 # ---------------------------------------------------------------------------
0025 # Helpers
0026 # ---------------------------------------------------------------------------
0027
0028 def extract_scope(dataset_name: str, strip_slash: bool = False):
0029     """
0030     Extract scope from a given dataset name.
0031
0032     Supports both formats:
0033     - Explicit colon format: scope:name (e.g., "user.pilot:dataset.name")
0034     - Inferred dot format:   scope.name (e.g., "user.pilot.dataset.name")
0035
0036     Based on the extract_scope method in rucio_comms/utils.py (RucioUtils).
0037
0038     Args:
0039         dataset_name: Dataset name in either format.
0040         strip_slash:  Whether to strip a trailing slash.
0041
0042     Returns:
0043         Tuple of (scope, name).
0044     """
0045     if strip_slash and dataset_name.endswith("/"):
0046         dataset_name = re.sub("/$", "", dataset_name)
0047
0048     # Handle explicit colon format: scope:name
0049     if ":" in dataset_name:
0050         parts = dataset_name.split(":", 1)
0051         if len(parts) == 2:
0052             return parts[0].strip(), parts[1].strip()
0053
0054     # Handle inferred dot format
0055     parts = dataset_name.split(".")
0056     if len(parts) < 2:
0057         raise ValueError(f"Dataset name must contain at least one dot or colon: {dataset_name}")
0058
0059     if dataset_name.startswith("user") or dataset_name.startswith("group"):
0060         if len(parts) >= 3:
0061             scope = ".".join(parts[0:2])
0062             name  = ".".join(parts[2:])
0063         else:
0064             scope = parts[0]
0065             name  = ".".join(parts[1:])
0066     else:
0067         scope = ".".join(parts[:-1])
0068         name  = parts[-1]
0069
0070     return scope, name
0071
0072
0073 def generate_vuid(scope: str, name: str) -> str:
0074     """
0075     Generate a Version UID (VUID) for a dataset.
0076
0077     Args:
0078         scope: Dataset scope.
0079         name:  Dataset name.
0080
0081     Returns:
0082         VUID string in UUID-like format.
0083     """
0084     vuid = hashlib.md5((scope + ":" + name).encode()).hexdigest()
0085     return f"{vuid[0:8]}-{vuid[8:12]}-{vuid[12:16]}-{vuid[16:20]}-{vuid[20:32]}"
0086
0087
0088 # ---------------------------------------------------------------------------
0089 # Checksum / file helpers
0090 # ---------------------------------------------------------------------------
0091
0092 def calculate_file_checksum(file_path: str, algorithm: str = 'md5', chunk_size: int = 4096) -> str:
0093     """
0094     Calculate the checksum of a file using the specified hash algorithm.
0095
0096     Args:
0097         file_path:  Path to the file.
0098         algorithm:  Hash algorithm name (e.g. ``'md5'``, ``'sha256'``).
0099         chunk_size: Size of chunks to read from the file.
0100
0101     Returns:
0102         Hex-encoded checksum string.
0103     """
0104     h = hashlib.new(algorithm)
0105     with open(file_path, 'rb') as f:
0106         while True:
0107             chunk = f.read(chunk_size)
0108             if not chunk:
0109                 break
0110             h.update(chunk)
0111     return h.hexdigest()
0112
0113
0114 def calculate_adler32_from_file(file_path, chunk_size=4096):
0115     """
0116     Calculates the Adler-32 checksum of a file.
0117
0118     Args:
0119         filepath (str): The path to the file.
0120         chunk_size (int): The size of chunks to read from the file.
0121
0122     Returns:
0123         int: The Adler-32 checksum of the file.
0124     """
0125     adler32_checksum = 1  # Initial Adler-32 value
0126
0127     try:
0128         with open(file_path, 'rb') as f:
0129             while True:
0130                 chunk = f.read(chunk_size)
0131                 if not chunk:
0132                     break
0133                 adler32_checksum = zlib.adler32(chunk, adler32_checksum)
0134         return adler32_checksum & 0xffffffff  # Ensure 32-bit unsigned result
0135     except Exception as e:
0136         raise OSError(f"Adler-32: problem with file {file_path}") from e
0137
0138 def register_file_on_rse(data_obj, file_path: str, file_name: str):
0139     """
0140     Register an uploaded file on RSE
0141
0142     This is a helper method to register a file on RSE after it has been uploaded.
0143
0144     It expects an object with some necessary attributes, e.g. the "data object" defined in relevant class.
0145     Attributes to be harvested from the "data object": client, rucio_did_client, replica_client, dataset, rse: str, scope: str
0146     """
0147
0148     adler = calculate_adler32_from_file(file_path)
0149     print(f"Adler32 checksum of the file {file_path}: {adler}")
0150
0151     try:
0152         # Step 1: Get file metadata
0153         file_size       = os.path.getsize(file_path)
0154         file_checksum   = calculate_file_checksum(file_path, 'md5')
0155
0156         print(f"File: {file_name}")
0157         print(f"Size: {file_size} bytes")
0158         print(f"MD5:  {file_checksum}")
0159
0160
0161         # Step 2: Check if DID already exists
0162         try:
0163             existing_did = data_obj.rucio_did_client.get_did(data_obj.rucio_scope, file_name)
0164             logger.info(f"DID already exists: {existing_did}")
0165         except DataIdentifierNotFound:
0166             # DID doesn't exist, we'll create it
0167             logger.info("DID doesn't exist yet, will create new one")
0168         except Exception as e:
0169             logger.warning(f"Unexpected error checking DID {data_obj.rucio_scope}:{file_name}: {e}")
0170
0171         dataset_folder = data_obj.dataset
0172
0173         # Register the replica
0174         data_obj.rucio_replica_client.add_replica(
0175             rse         = data_obj.rse,
0176             scope       = data_obj.rucio_scope,
0177             name        = file_name,
0178             bytes_      = file_size,
0179             adler32     = f'{adler:x}',
0180             pfn         = f'root://dcintdoor.sdcc.bnl.gov:1094/pnfs/sdcc.bnl.gov/eic/epic/disk/swfdaqtest/{dataset_folder}/{file_name}'
0181             )
0182
0183         print(f"✓ Replica registered on RSE: {data_obj.rse}")
0184
0185         return True
0186
0187     except RSENotFound:
0188         print(f"✗ Error: RSE '{data_obj.rse}' not found")
0189         return False
0190     except Exception as e:
0191         print(f"✗ Error registering file: {str(e)}")
0192         return False
0193
0194
0195 # ---------------------------------------------------------------------------
0196 # Dataset operations  (standalone equivalents of DatasetManager methods)
0197 # ---------------------------------------------------------------------------
0198
0199 def create_dataset(dataset_name: str, lifetime_days: Optional[int] = None,
0200                    open_dataset: bool = True, client=None):
0201     """
0202     Create a Rucio dataset.
0203
0204     Standalone equivalent of ``DatasetManager.create_dataset``.
0205
0206     Args:
0207         dataset_name:  Full dataset identifier (``scope:name`` or dot-separated).
0208         lifetime_days: Optional lifetime in days.
0209         open_dataset:  Whether the dataset should be left open (default True).
0210         client:        An existing ``rucio.client.Client`` instance.
0211                        A new one is created when *None*.
0212
0213     Returns:
0214         dict with keys ``scope``, ``name``, ``duid`` on success, or *None* on
0215         failure.
0216     """
0217     if client is None:
0218         client = RucioClient()
0219
0220     try:
0221         scope, name = extract_scope(dataset_name)
0222         logger.info(f"Creating dataset: {scope}:{name}")
0223
0224         # Build metadata
0225         meta = {}
0226         if lifetime_days is not None:
0227             meta['lifetime'] = lifetime_days * 86400  # seconds
0228
0229         # Create the dataset DID
0230         try:
0231             client.add_dataset(scope=scope, name=name, meta=meta,
0232                                lifetime=meta.get('lifetime'))
0233             logger.info(f"Dataset created: {scope}:{name}")
0234         except DataIdentifierAlreadyExists:
0235             logger.info(f"Dataset already exists: {scope}:{name}")
0236             # Apply lifetime to existing dataset if requested
0237             if lifetime_days is not None:
0238                 client.set_metadata(scope=scope, name=name,
0239                                     key='lifetime',
0240                                     value=lifetime_days * 86400)
0241                 logger.info(f"Updated lifetime for existing dataset: "
0242                             f"{scope}:{name} -> {lifetime_days} days")
0243
0244         # Set open/closed status
0245         if open_dataset:
0246             try:
0247                 client.set_status(scope=scope, name=name, open=True)
0248             except Exception:
0249                 pass  # may already be open
0250
0251         # Generate identifiers
0252         vuid = generate_vuid(scope, name)
0253         duid = vuid
0254         result = {
0255             'scope': scope,
0256             'name':  name,
0257             'duid':  duid,
0258             'vuid':  vuid,
0259         }
0260         logger.info(f"Dataset ready: {scope}:{name}  duid={duid}")
0261         return result
0262
0263     except Exception as e:
0264         logger.error(f"Failed to create dataset {dataset_name}: {e}")
0265         return None
0266
0267
0268 # ---------------------------------------------------------------------------
0269 # File-to-dataset attachment  (standalone equivalent of FileManager method)
0270 # ---------------------------------------------------------------------------
0271
0272 def add_files_to_dataset(files, dataset_name: str,
0273                          dataset_scope: Optional[str] = None, rse: Optional[str] = None,
0274                          client=None):
0275     """
0276     Add files to a Rucio dataset.
0277
0278     Standalone equivalent of ``FileManager.add_files_to_dataset``.
0279
0280     Args:
0281         files:         List of LFN strings (``scope:name`` or dot-separated).
0282         dataset_name:  Target dataset name (``scope:name`` or dot-separated).
0283         dataset_scope: Explicit dataset scope (extracted from *dataset_name*
0284                        when *None*).
0285         rse:           Optional RSE constraint.
0286         client:        An existing ``rucio.client.Client`` instance.
0287                        A new one is created when *None*.
0288
0289     Returns:
0290         True on success.
0291
0292     Raises:
0293         RuntimeError: If the operation fails.
0294     """
0295     if client is None:
0296         client = RucioClient()
0297
0298     try:
0299         if dataset_scope is None:
0300             dataset_scope, dataset_name = extract_scope(dataset_name)
0301
0302         logger.info(f"Adding {len(files)} file(s) to dataset: "
0303                     f"{dataset_scope}:{dataset_name}")
0304
0305         # Build Rucio file dicts
0306         file_dicts = []
0307         for item in files:
0308             if isinstance(item, str):
0309                 file_scope, lfn = extract_scope(item)
0310                 file_dicts.append({'scope': file_scope, 'name': lfn})
0311             else:
0312                 raise ValueError(f"Invalid file item type: {type(item)}")
0313
0314         # Attach in batches of 1000
0315         batch_size = 1000
0316         for i in range(0, len(file_dicts), batch_size):
0317             batch = file_dicts[i:i + batch_size]
0318             try:
0319                 client.add_files_to_dataset(
0320                     scope=dataset_scope, name=dataset_name,
0321                     files=batch, rse=rse,
0322                 )
0323                 logger.debug(f"Added batch of {len(batch)} file(s) to dataset")
0324             except FileAlreadyExists:
0325                 # Retry individually so we skip only true duplicates
0326                 for fd in batch:
0327                     try:
0328                         client.add_files_to_dataset(
0329                             scope=dataset_scope, name=dataset_name,
0330                             files=[fd], rse=rse,
0331                         )
0332                     except FileAlreadyExists:
0333                         logger.debug(f"File already in dataset: {fd['name']}")
0334
0335         logger.info(f"Successfully added files to dataset: "
0336                     f"{dataset_scope}:{dataset_name}")
0337         return True
0338
0339     except Exception as e:
0340         error_msg = (f"Failed to add files to dataset "
0341                      f"{dataset_scope}:{dataset_name}: {e}")
0342         logger.error(error_msg)
0343         raise RuntimeError(error_msg) from e