File indexing completed on 2026-09-01 09:34:30
0001 from abc import ABC, abstractmethod
0002 from typing import Any
0003
0004 from .models import FileDID
0005
0006
0007 class DatasetCatalog(ABC):
0008 """Minimal interface needed by the decision box.
0009
0010 Production prompt processing uses the Rucio-backed implementation below.
0011 """
0012
0013 @abstractmethod
0014 def ensure_dataset(self, dataset_did: str, open_dataset: bool = True) -> None:
0015 """Create the dataset if needed and keep its open/closed state."""
0016
0017 @abstractmethod
0018 def attach_file(self, dataset_did: str, file_did: FileDID) -> bool:
0019 """Attach a file DID. Return True if this call added a new member."""
0020
0021 @abstractmethod
0022 def close_dataset(self, dataset_did: str) -> None:
0023 """Mark a dataset closed."""
0024
0025
0026 class DatasetStateLookupError(RuntimeError):
0027 """Raised when a dataset open/closed state cannot be determined."""
0028
0029
0030 class RucioDatasetCatalog(DatasetCatalog):
0031 """Rucio-backed catalog adapter.
0032
0033 This adapter follows the same helper APIs used by the existing data agent
0034 and reuses the data agent's initialized Rucio client.
0035 """
0036
0037 def __init__(self, client, lifetime_days: int | None = None):
0038 self.client = client
0039 self.lifetime_days = lifetime_days
0040 self._use_rucio_utils = False
0041 self._closed_datasets: set[str] = set()
0042 try:
0043 from swf_common_lib.rucio_utils import create_dataset, add_files_to_dataset
0044
0045 self._create_dataset = create_dataset
0046 self._add_files_to_dataset = add_files_to_dataset
0047 self._use_rucio_utils = True
0048 except ModuleNotFoundError as exc:
0049 if exc.name not in {"swf_common_lib", "swf_common_lib.rucio_utils"}:
0050 raise
0051 from rucio.common.exception import DataIdentifierAlreadyExists
0052 from rucio_comms import DatasetManager, FileManager
0053
0054 self.dataset_manager = DatasetManager()
0055 self.file_manager = FileManager(rucio_client=self.client)
0056 self.data_identifier_already_exists = DataIdentifierAlreadyExists
0057
0058 def ensure_dataset(self, dataset_did: str, open_dataset: bool = True) -> None:
0059 if open_dataset and self._dataset_is_closed(dataset_did):
0060 raise ValueError(f"dataset {dataset_did} is already closed")
0061 if self._use_rucio_utils:
0062 result = self._create_dataset(
0063 dataset_name=dataset_did,
0064 lifetime_days=self.lifetime_days,
0065 open_dataset=open_dataset,
0066 client=self.client,
0067 )
0068 if not result:
0069 raise RuntimeError(f"failed to create dataset {dataset_did}")
0070 return
0071 try:
0072 self.dataset_manager.create_dataset(
0073 dataset_name=dataset_did,
0074 lifetime_days=self.lifetime_days,
0075 open_dataset=open_dataset,
0076 )
0077 except self.data_identifier_already_exists:
0078 return
0079
0080 def attach_file(self, dataset_did: str, file_did: FileDID) -> bool:
0081 if self._dataset_is_closed(dataset_did):
0082 raise ValueError(f"dataset {dataset_did} is closed")
0083 self.ensure_dataset(dataset_did, open_dataset=True)
0084 if self._use_rucio_utils:
0085 result = self._add_files_to_dataset([str(file_did)], dataset_did, client=self.client)
0086 else:
0087 result = self.file_manager.add_files_to_dataset([str(file_did)], dataset_did)
0088 return bool(result)
0089
0090 def close_dataset(self, dataset_did: str) -> None:
0091 scope, name = dataset_did.split(":", 1)
0092 self.client.set_status(scope=scope, name=name, open=False)
0093 self._closed_datasets.add(dataset_did)
0094
0095 def _dataset_is_closed(self, dataset_did: str) -> bool:
0096 if dataset_did in self._closed_datasets:
0097 return True
0098 scope, name = dataset_did.split(":", 1)
0099 for method_name in ("get_did", "get_metadata"):
0100 method = getattr(self.client, method_name, None)
0101 if method is None:
0102 continue
0103 try:
0104 metadata = method(scope=scope, name=name)
0105 except Exception as exc:
0106 if self._is_did_not_found(exc):
0107 continue
0108 raise DatasetStateLookupError(f"failed to look up dataset {dataset_did}") from exc
0109 is_open = self._metadata_open_state(metadata)
0110 if is_open is None:
0111 continue
0112 if not is_open:
0113 self._closed_datasets.add(dataset_did)
0114 return True
0115 return False
0116 return False
0117
0118 @staticmethod
0119 def _is_did_not_found(exc: Exception) -> bool:
0120 return any(
0121 cls.__name__ in {"DataIdentifierNotFound", "DIDNotFound", "DataIdentifierNotFoundError"}
0122 for cls in type(exc).__mro__
0123 )
0124
0125 @staticmethod
0126 def _metadata_open_state(metadata: Any) -> bool | None:
0127 if not metadata:
0128 return None
0129 for key in ("is_open", "open"):
0130 if key in metadata:
0131 value = metadata[key]
0132 if isinstance(value, str):
0133 return value.strip().lower() not in {"0", "false", "no", "off"}
0134 return bool(value)
0135 return None
