Coverage for python / lsst / daf / butler / datastores / file_datastore / retrieve_artifacts.py: 28%
136 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-04 02:35 -0800
1# This file is part of daf_butler.
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = ("ZipIndex", "determine_destination_for_retrieved_artifact", "retrieve_and_zip", "unpack_zips")
32import logging
33import tempfile
34import uuid
35import zipfile
36from collections.abc import Iterable
37from typing import ClassVar, Literal, Protocol, Self
39from pydantic import BaseModel
41from lsst.daf.butler import (
42 DatasetIdFactory,
43 DatasetRef,
44 SerializedDatasetRefContainers,
45 SerializedDatasetRefContainerV1,
46)
47from lsst.daf.butler.datastore.stored_file_info import SerializedStoredFileInfo
48from lsst.resources import ResourcePath, ResourcePathExpression
50_LOG = logging.getLogger(__name__)
53class ArtifactIndexInfo(BaseModel):
54 """Information related to an artifact in an index."""
56 info: SerializedStoredFileInfo
57 """Datastore record information for this file artifact."""
59 ids: set[uuid.UUID]
60 """Dataset IDs for this artifact."""
62 def append(self, id_: uuid.UUID) -> None:
63 """Add an additional dataset ID.
65 Parameters
66 ----------
67 id_ : `uuid.UUID`
68 Additional dataset ID to associate with this artifact.
69 """
70 self.ids.add(id_)
72 @classmethod
73 def from_single(cls, info: SerializedStoredFileInfo, id_: uuid.UUID) -> Self:
74 """Create a mapping from a single ID and info.
76 Parameters
77 ----------
78 info : `SerializedStoredFileInfo`
79 Datastore record for this artifact.
80 id_ : `uuid.UUID`
81 First dataset ID to associate with this artifact.
82 """
83 return cls(info=info, ids=[id_])
85 def subset(self, ids: Iterable[uuid.UUID]) -> Self:
86 """Replace the IDs with a subset of the IDs and return a new instance.
88 Parameters
89 ----------
90 ids : `~collections.abc.Iterable` [ `uuid.UUID` ]
91 Subset of IDs to keep.
93 Returns
94 -------
95 subsetted : `ArtifactIndexInfo`
96 New instance with the requested subset.
98 Raises
99 ------
100 ValueError
101 Raised if the given IDs is not a subset of the current IDs.
102 """
103 subset = set(ids)
104 if subset - self.ids:
105 raise ValueError(f"Given subset of {subset} is not a subset of {self.ids}")
106 return type(self)(ids=subset, info=self.info.model_copy())
109class ZipIndex(BaseModel):
110 """Index of a Zip file of Butler datasets.
112 A file can be associated with multiple butler datasets and a single
113 butler dataset can be associated with multiple files. This model
114 provides the necessary information for ingesting back into a Butler
115 file datastore.
116 """
118 index_version: Literal["V1"] = "V1"
120 refs: SerializedDatasetRefContainers
121 """Deduplicated information for all the `DatasetRef` in the index."""
123 artifact_map: dict[str, ArtifactIndexInfo]
124 """Mapping of each Zip member to associated lookup information."""
126 index_name: ClassVar[str] = "_butler_zip_index.json"
127 """Name to use when saving the index to a file."""
129 def generate_uuid5(self) -> uuid.UUID:
130 """Create a UUID based on the Zip index.
132 Returns
133 -------
134 id_ : `uuid.UUID`
135 A UUID5 created from the paths inside the Zip file. Guarantees
136 that if the Zip file is regenerated with exactly the same file
137 paths the same answer will be returned.
138 """
139 # Options are:
140 # - uuid5 based on file paths in zip
141 # - uuid5 based on ref uuids.
142 # - checksum derived from the above.
143 # - uuid5 from file paths and dataset refs.
144 # Do not attempt to include file contents in UUID.
145 # Start with uuid5 from file paths.
146 data = ",".join(sorted(self.artifact_map.keys()))
147 # No need to come up with a different namespace.
148 return uuid.uuid5(DatasetIdFactory.NS_UUID, data)
150 def __len__(self) -> int:
151 """Return the number of files in the Zip."""
152 return len(self.artifact_map)
154 def calculate_zip_file_name(self) -> str:
155 """Calculate the default name for the Zip file based on the index
156 contents.
158 Returns
159 -------
160 name : `str`
161 Name of the zip file based on index.
162 """
163 return f"{self.generate_uuid5()}.zip"
165 def calculate_zip_file_path_in_store(self) -> str:
166 """Calculate the relative path inside a datastore that should be
167 used for this Zip file.
169 Returns
170 -------
171 path_in_store : `str`
172 Relative path to use for Zip file in datastore.
173 """
174 zip_name = self.calculate_zip_file_name()
175 return f"zips/{zip_name[:4]}/{zip_name}"
177 def write_index(self, dir: ResourcePath) -> ResourcePath:
178 """Write the index to the specified directory.
180 Parameters
181 ----------
182 dir : `~lsst.resources.ResourcePath`
183 Directory to write the index file to.
185 Returns
186 -------
187 index_path : `~lsst.resources.ResourcePath`
188 Path to the index file that was written.
189 """
190 index_path = dir.join(self.index_name, forceDirectory=False)
191 with index_path.open("w") as fd:
192 # Need to include unset/default values so that the version
193 # discriminator field for refs container appears in the
194 # serialization.
195 print(self.model_dump_json(), file=fd)
196 return index_path
198 @classmethod
199 def calc_relative_paths(
200 cls, root: ResourcePath, paths: Iterable[ResourcePath]
201 ) -> dict[ResourcePath, str]:
202 """Calculate the path to use inside the Zip file from the full path.
204 Parameters
205 ----------
206 root : `lsst.resources.ResourcePath`
207 The reference root directory.
208 paths : `~collections.abc.Iterable` [ `lsst.resources.ResourcePath` ]
209 The paths to the files that should be included in the Zip file.
211 Returns
212 -------
213 abs_to_rel : `dict` [ `~lsst.resources.ResourcePath`, `str` ]
214 Mapping of the original file path to the relative path to use
215 in Zip file.
216 """
217 file_to_relative: dict[ResourcePath, str] = {}
218 for p in paths:
219 # It is an error if there is no relative path.
220 rel = p.relative_to(root)
221 if rel is None:
222 raise RuntimeError(f"Unexepectedly unable to calculate relative path of {p} to {root}.")
223 file_to_relative[p] = rel
224 return file_to_relative
226 @classmethod
227 def from_artifact_map(
228 cls,
229 refs: Iterable[DatasetRef],
230 artifact_map: dict[ResourcePath, ArtifactIndexInfo],
231 root: ResourcePath,
232 ) -> Self:
233 """Create an index from the mappings returned from
234 `Datastore.retrieveArtifacts`.
236 Parameters
237 ----------
238 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
239 Datasets present in the index.
240 artifact_map : `dict` [ `lsst.resources.ResourcePath`, `ArtifactMap` ]
241 Mapping of artifact path to information linking it to the
242 associated refs and datastore information.
243 root : `lsst.resources.ResourcePath`
244 Root path to be removed from all the paths before creating the
245 index.
246 """
247 if not refs:
248 return cls(
249 refs=SerializedDatasetRefContainerV1.from_refs(refs),
250 artifact_map={},
251 )
253 # Calculate the paths relative to the given root since the Zip file
254 # uses relative paths.
255 file_to_relative = cls.calc_relative_paths(root, artifact_map.keys())
257 simplified_refs = SerializedDatasetRefContainerV1.from_refs(refs)
259 # Convert the artifact mapping to relative path.
260 relative_artifact_map = {file_to_relative[path]: info for path, info in artifact_map.items()}
262 return cls(
263 refs=simplified_refs,
264 artifact_map=relative_artifact_map,
265 )
267 @classmethod
268 def from_zip_file(cls, zip_path: ResourcePath) -> Self:
269 """Given a path to a Zip file return the index.
271 Parameters
272 ----------
273 zip_path : `lsst.resources.ResourcePath`
274 Path to the Zip file.
275 """
276 with zip_path.open("rb") as fd, zipfile.ZipFile(fd) as zf:
277 return cls.from_open_zip(zf)
279 @classmethod
280 def from_open_zip(cls, zf: zipfile.ZipFile) -> Self:
281 json_data = zf.read(cls.index_name)
282 return cls.model_validate_json(json_data)
285def determine_destination_for_retrieved_artifact(
286 destination_directory: ResourcePath, source_path: ResourcePath, preserve_path: bool, prefix: str = ""
287) -> ResourcePath:
288 """Determine destination path for an artifact retrieved from a datastore.
290 Parameters
291 ----------
292 destination_directory : `ResourcePath`
293 Path to the output directory where file will be stored.
294 source_path : `ResourcePath`
295 Path to the source file to be transferred. This may be relative to the
296 datastore root, or an absolute path.
297 preserve_path : `bool`
298 If `True` the full path of the artifact within the datastore
299 is preserved. If `False` the final file component of the path
300 is used.
301 prefix : `str`, optional
302 Prefix to add to the file name if ``preserve_path`` is `False`.
304 Returns
305 -------
306 destination_uri : `~lsst.resources.ResourcePath`
307 Absolute path to the output destination.
308 """
309 destination_directory = destination_directory.abspath()
311 target_path: ResourcePathExpression
312 if preserve_path:
313 target_path = source_path
314 if target_path.isabs():
315 # This is an absolute path to an external file.
316 # Use the full path.
317 target_path = target_path.relativeToPathRoot
318 else:
319 target_path = source_path.basename()
320 if prefix:
321 target_path = prefix + target_path
323 target_uri = destination_directory.join(target_path).abspath()
324 if target_uri.relative_to(destination_directory) is None:
325 raise ValueError(f"File path attempts to escape destination directory: '{source_path}'")
326 return target_uri
329class RetrievalCallable(Protocol):
330 def __call__( 330 ↛ exitline 330 didn't return from function '__call__' because
331 self,
332 refs: Iterable[DatasetRef],
333 destination: ResourcePath,
334 transfer: str,
335 preserve_path: bool,
336 overwrite: bool,
337 write_index: bool,
338 add_prefix: bool,
339 ) -> dict[ResourcePath, ArtifactIndexInfo]: ...
342def retrieve_and_zip(
343 refs: Iterable[DatasetRef],
344 destination: ResourcePathExpression,
345 retrieval_callback: RetrievalCallable,
346 overwrite: bool = True,
347) -> ResourcePath:
348 """Retrieve artifacts from a Butler and place in ZIP file.
350 Parameters
351 ----------
352 refs : `collections.abc.Iterable` [ `DatasetRef` ]
353 The datasets to be included in the zip file. Must all be from
354 the same dataset type.
355 destination : `lsst.resources.ResourcePath`
356 Directory to write the new ZIP file. This directory will
357 also be used as a staging area for the datasets being downloaded
358 from the datastore.
359 retrieval_callback : `~collections.abc.Callable`
360 Bound method for a function that can retrieve the artifacts and
361 return the metadata necessary for creating the zip index. For example
362 `lsst.daf.butler.datastore.Datastore.retrieveArtifacts`.
363 overwrite : `bool`, optional
364 If `False` the output Zip will not be written if a file of the
365 same name is already present in ``destination``.
367 Returns
368 -------
369 zip_file : `lsst.resources.ResourcePath`
370 The path to the new ZIP file.
372 Raises
373 ------
374 ValueError
375 Raised if there are no refs to retrieve.
376 """
377 if not refs:
378 raise ValueError("Requested Zip file with no contents.")
380 outdir = ResourcePath(destination, forceDirectory=True)
381 if not outdir.isdir():
382 raise ValueError(f"Destination location must refer to a directory. Given {destination}")
384 if not outdir.exists():
385 outdir.mkdir()
387 # Simplest approach:
388 # - create temp dir in destination
389 # - Run retrieveArtifacts to that temp dir
390 # - Create zip file with unique name.
391 # - Delete temp dir
392 # - Add index file to ZIP.
393 # - Return name of zip file.
394 with tempfile.TemporaryDirectory(dir=outdir.ospath, ignore_cleanup_errors=True) as tmpdir:
395 tmpdir_path = ResourcePath(tmpdir, forceDirectory=True)
396 # Retrieve the artifacts and write the index file. Strip paths.
397 artifact_map = retrieval_callback(
398 refs=refs,
399 destination=tmpdir_path,
400 transfer="auto",
401 preserve_path=False,
402 overwrite=False,
403 write_index=True,
404 add_prefix=True,
405 )
407 # Read the index to construct file name.
408 index_path = tmpdir_path.join(ZipIndex.index_name, forceDirectory=False)
409 index_json = index_path.read()
410 index = ZipIndex.model_validate_json(index_json)
412 # Use unique name based on files in Zip.
413 zip_file_name = index.calculate_zip_file_name()
414 zip_path = outdir.join(zip_file_name, forceDirectory=False)
415 if not overwrite and zip_path.exists():
416 raise FileExistsError(f"Output Zip at {zip_path} already exists but cannot overwrite.")
417 with zipfile.ZipFile(zip_path.ospath, "w") as zip:
418 zip.write(index_path.ospath, index_path.basename(), compress_type=zipfile.ZIP_DEFLATED)
419 for path, name in index.calc_relative_paths(tmpdir_path, list(artifact_map)).items():
420 zip.write(path.ospath, name)
422 return zip_path
425def unpack_zips(
426 zips_to_transfer: Iterable[ResourcePath],
427 allowed_ids: set[uuid.UUID],
428 destination: ResourcePath,
429 preserve_path: bool,
430 overwrite: bool,
431) -> dict[ResourcePath, ArtifactIndexInfo]:
432 """Transfer the Zip files and unpack them in the destination directory.
434 Parameters
435 ----------
436 zips_to_transfer : `~collections.abc.Iterable` \
437 [ `~lsst.resources.ResourcePath` ]
438 Paths to Zip files to unpack. These must be Zip files that include
439 the index information and were created by the Butler.
440 allowed_ids : `set` [ `uuid.UUID` ]
441 All the possible dataset IDs for which artifacts should be extracted
442 from the Zip file. If an ID in the Zip file is not present in this
443 list the artifact will not be extracted from the Zip.
444 destination : `~lsst.resources.ResourcePath`
445 Output destination for the Zip contents.
446 preserve_path : `bool`
447 Whether to include subdirectories during extraction. If `True` a
448 directory will be made per Zip.
449 overwrite : `bool`, optional
450 If `True` allow transfers to overwrite existing files at the
451 destination.
453 Returns
454 -------
455 artifact_map : `dict` \
456 [ `~lsst.resources.ResourcePath`, `ArtifactIndexInfo` ]
457 Path linking Zip contents location to associated artifact information.
458 """
459 artifact_map: dict[ResourcePath, ArtifactIndexInfo] = {}
460 for source_uri in zips_to_transfer:
461 _LOG.debug("Unpacking zip file %s", source_uri)
462 # Assume that downloading to temporary location is more efficient
463 # than trying to read the contents remotely.
464 with ResourcePath.temporary_uri(suffix=".zip") as temp:
465 temp.transfer_from(source_uri, transfer="auto")
466 index = ZipIndex.from_zip_file(temp)
468 if preserve_path:
469 subdir = ResourcePath(
470 index.calculate_zip_file_path_in_store(), forceDirectory=False, forceAbsolute=False
471 ).dirname()
472 outdir = destination.join(subdir)
473 else:
474 outdir = destination
475 outdir.mkdir()
476 with temp.open("rb") as fd, zipfile.ZipFile(fd) as zf:
477 for path_in_zip, artifact_info in index.artifact_map.items():
478 # Skip if this specific dataset ref is not requested.
479 included_ids = artifact_info.ids & allowed_ids
480 if included_ids:
481 # Do not apply a new prefix since the zip file
482 # should already have a prefix.
483 output_path = outdir.join(path_in_zip, forceDirectory=False)
484 if not overwrite and output_path.exists():
485 raise FileExistsError(
486 f"Destination path '{output_path}' already exists. "
487 "Extraction from Zip cannot be completed."
488 )
489 zf.extract(path_in_zip, path=outdir.ospath)
490 artifact_map[output_path] = artifact_info.subset(included_ids)
491 return artifact_map