Coverage for python / lsst / rucio / register / export.py: 0%

1# This file is part of rucio_register 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22 

23from typing import Any 

24 

25import click 

26 

27from lsst.daf.butler import Butler, CollectionType, FileDataset 

28from lsst.daf.butler.cli.opt import query_datasets_options 

29 

30 

31@click.command(short_help="Export registry information about datasets.") 

32@click.option("--root", help="URI root for existing direct ingests to be stripped.") 

33@click.option("--filename", default="export.yaml", help="Output filename (default=export.yaml).") 

34@query_datasets_options(repo=True, useArguments=True, use_order_by=False, showUri=False) 

35def main(repo: str, glob: Any, filename: str, root: str | None = None, **kwargs: Any) -> None: 

36 """Export the registry information about datasets selected by collection, 

37 dataset type, and where clause to an export.yaml file. 

38 """ 

39 if "collections" not in kwargs: 

40 raise RuntimeError("Collection(s) option is required.") 

41 

42 butler = Butler(repo) 

43 ds_types = set([ds.name for ds in butler.registry.queryDatasetTypes(glob)]) 

44 

45 if root is not None: 

46 if not root.endswith("/"): 

47 root += "/" 

48 

49 def rewrite_file_dataset(file_dataset: FileDataset) -> FileDataset: 

50 file_dataset.path = file_dataset.path.removeprefix(root) 

51 return file_dataset 

52 

53 rewrite = rewrite_file_dataset 

54 print(f"Stripping root prefix: {root}") 

55 else: 

56 rewrite = None 

57 

58 print(f"Output to: {filename}") 

59 with butler.export(filename=filename, format="yaml", transfer=None) as export: 

60 actual_ds_types = set() 

61 

62 # We only save calibration and run collections, not complete chains. 

63 # This is expected to be more useful for DRP. 

64 for collection_info in butler.collections.query_info( 

65 kwargs["collections"], 

66 flatten_chains=True, 

67 include_summary=True, 

68 summary_datasets=ds_types, 

69 ): 

70 present_ds_types = ds_types & collection_info.dataset_types 

71 if present_ds_types: 

72 if collection_info.type == CollectionType.CALIBRATION: 

73 print(f"Saving collection associations: {collection_info.name}") 

74 export.saveCollection(collection_info.name) 

75 actual_ds_types |= present_ds_types 

76 

77 for ds_type in actual_ds_types: 

78 print(f"Saving dataset type: {ds_type}") 

79 export.saveDatasets( 

80 butler.query_datasets(ds_type, with_dimension_records=True, **kwargs), 

81 rewrite=rewrite, 

82 )