Coverage for python / lsst / daf / butler / datastores / file_datastore / retrieve_artifacts.py: 27%
133 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-10 01:37 -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 json_data = zf.read(cls.index_name)
278 return cls.model_validate_json(json_data)
281def determine_destination_for_retrieved_artifact(
282 destination_directory: ResourcePath, source_path: ResourcePath, preserve_path: bool, prefix: str = ""
283) -> ResourcePath:
284 """Determine destination path for an artifact retrieved from a datastore.
286 Parameters
287 ----------
288 destination_directory : `ResourcePath`
289 Path to the output directory where file will be stored.
290 source_path : `ResourcePath`
291 Path to the source file to be transferred. This may be relative to the
292 datastore root, or an absolute path.
293 preserve_path : `bool`
294 If `True` the full path of the artifact within the datastore
295 is preserved. If `False` the final file component of the path
296 is used.
297 prefix : `str`, optional
298 Prefix to add to the file name if ``preserve_path`` is `False`.
300 Returns
301 -------
302 destination_uri : `~lsst.resources.ResourcePath`
303 Absolute path to the output destination.
304 """
305 destination_directory = destination_directory.abspath()
307 target_path: ResourcePathExpression
308 if preserve_path:
309 target_path = source_path
310 if target_path.isabs():
311 # This is an absolute path to an external file.
312 # Use the full path.
313 target_path = target_path.relativeToPathRoot
314 else:
315 target_path = source_path.basename()
316 if prefix:
317 target_path = prefix + target_path
319 target_uri = destination_directory.join(target_path).abspath()
320 if target_uri.relative_to(destination_directory) is None:
321 raise ValueError(f"File path attempts to escape destination directory: '{source_path}'")
322 return target_uri
325class RetrievalCallable(Protocol):
326 def __call__( 326 ↛ exitline 326 didn't return from function '__call__' because
327 self,
328 refs: Iterable[DatasetRef],
329 destination: ResourcePath,
330 transfer: str,
331 preserve_path: bool,
332 overwrite: bool,
333 write_index: bool,
334 add_prefix: bool,
335 ) -> dict[ResourcePath, ArtifactIndexInfo]: ...
338def retrieve_and_zip(
339 refs: Iterable[DatasetRef],
340 destination: ResourcePathExpression,
341 retrieval_callback: RetrievalCallable,
342 overwrite: bool = True,
343) -> ResourcePath:
344 """Retrieve artifacts from a Butler and place in ZIP file.
346 Parameters
347 ----------
348 refs : `collections.abc.Iterable` [ `DatasetRef` ]
349 The datasets to be included in the zip file. Must all be from
350 the same dataset type.
351 destination : `lsst.resources.ResourcePath`
352 Directory to write the new ZIP file. This directory will
353 also be used as a staging area for the datasets being downloaded
354 from the datastore.
355 retrieval_callback : `~collections.abc.Callable`
356 Bound method for a function that can retrieve the artifacts and
357 return the metadata necessary for creating the zip index. For example
358 `lsst.daf.butler.datastore.Datastore.retrieveArtifacts`.
359 overwrite : `bool`, optional
360 If `False` the output Zip will not be written if a file of the
361 same name is already present in ``destination``.
363 Returns
364 -------
365 zip_file : `lsst.resources.ResourcePath`
366 The path to the new ZIP file.
368 Raises
369 ------
370 ValueError
371 Raised if there are no refs to retrieve.
372 """
373 if not refs:
374 raise ValueError("Requested Zip file with no contents.")
376 outdir = ResourcePath(destination, forceDirectory=True)
377 if not outdir.isdir():
378 raise ValueError(f"Destination location must refer to a directory. Given {destination}")
380 if not outdir.exists():
381 outdir.mkdir()
383 # Simplest approach:
384 # - create temp dir in destination
385 # - Run retrieveArtifacts to that temp dir
386 # - Create zip file with unique name.
387 # - Delete temp dir
388 # - Add index file to ZIP.
389 # - Return name of zip file.
390 with tempfile.TemporaryDirectory(dir=outdir.ospath, ignore_cleanup_errors=True) as tmpdir:
391 tmpdir_path = ResourcePath(tmpdir, forceDirectory=True)
392 # Retrieve the artifacts and write the index file. Strip paths.
393 artifact_map = retrieval_callback(
394 refs=refs,
395 destination=tmpdir_path,
396 transfer="auto",
397 preserve_path=False,
398 overwrite=False,
399 write_index=True,
400 add_prefix=True,
401 )
403 # Read the index to construct file name.
404 index_path = tmpdir_path.join(ZipIndex.index_name, forceDirectory=False)
405 index_json = index_path.read()
406 index = ZipIndex.model_validate_json(index_json)
408 # Use unique name based on files in Zip.
409 zip_file_name = index.calculate_zip_file_name()
410 zip_path = outdir.join(zip_file_name, forceDirectory=False)
411 if not overwrite and zip_path.exists():
412 raise FileExistsError(f"Output Zip at {zip_path} already exists but cannot overwrite.")
413 with zipfile.ZipFile(zip_path.ospath, "w") as zip:
414 zip.write(index_path.ospath, index_path.basename(), compress_type=zipfile.ZIP_DEFLATED)
415 for path, name in index.calc_relative_paths(tmpdir_path, list(artifact_map)).items():
416 zip.write(path.ospath, name)
418 return zip_path
421def unpack_zips(
422 zips_to_transfer: Iterable[ResourcePath],
423 allowed_ids: set[uuid.UUID],
424 destination: ResourcePath,
425 preserve_path: bool,
426 overwrite: bool,
427) -> dict[ResourcePath, ArtifactIndexInfo]:
428 """Transfer the Zip files and unpack them in the destination directory.
430 Parameters
431 ----------
432 zips_to_transfer : `~collections.abc.Iterable` \
433 [ `~lsst.resources.ResourcePath` ]
434 Paths to Zip files to unpack. These must be Zip files that include
435 the index information and were created by the Butler.
436 allowed_ids : `set` [ `uuid.UUID` ]
437 All the possible dataset IDs for which artifacts should be extracted
438 from the Zip file. If an ID in the Zip file is not present in this
439 list the artifact will not be extracted from the Zip.
440 destination : `~lsst.resources.ResourcePath`
441 Output destination for the Zip contents.
442 preserve_path : `bool`
443 Whether to include subdirectories during extraction. If `True` a
444 directory will be made per Zip.
445 overwrite : `bool`, optional
446 If `True` allow transfers to overwrite existing files at the
447 destination.
449 Returns
450 -------
451 artifact_map : `dict` \
452 [ `~lsst.resources.ResourcePath`, `ArtifactIndexInfo` ]
453 Path linking Zip contents location to associated artifact information.
454 """
455 artifact_map: dict[ResourcePath, ArtifactIndexInfo] = {}
456 for source_uri in zips_to_transfer:
457 _LOG.debug("Unpacking zip file %s", source_uri)
458 # Assume that downloading to temporary location is more efficient
459 # than trying to read the contents remotely.
460 with ResourcePath.temporary_uri(suffix=".zip") as temp:
461 temp.transfer_from(source_uri, transfer="auto")
462 index = ZipIndex.from_zip_file(temp)
464 if preserve_path:
465 subdir = ResourcePath(
466 index.calculate_zip_file_path_in_store(), forceDirectory=False, forceAbsolute=False
467 ).dirname()
468 outdir = destination.join(subdir)
469 else:
470 outdir = destination
471 outdir.mkdir()
472 with temp.open("rb") as fd, zipfile.ZipFile(fd) as zf:
473 for path_in_zip, artifact_info in index.artifact_map.items():
474 # Skip if this specific dataset ref is not requested.
475 included_ids = artifact_info.ids & allowed_ids
476 if included_ids:
477 # Do not apply a new prefix since the zip file
478 # should already have a prefix.
479 output_path = outdir.join(path_in_zip, forceDirectory=False)
480 if not overwrite and output_path.exists():
481 raise FileExistsError(
482 f"Destination path '{output_path}' already exists. "
483 "Extraction from Zip cannot be completed."
484 )
485 zf.extract(path_in_zip, path=outdir.ospath)
486 artifact_map[output_path] = artifact_info.subset(included_ids)
487 return artifact_map