From 67df37b1b59fd69468af1f0e3e377cd8e842faf7 Mon Sep 17 00:00:00 2001 From: Andy Salnikov Date: Tue, 13 Aug 2024 14:41:56 -0700 Subject: [PATCH] Update RepoExportContext to use new query system instead of queryDatasetAssociations. --- python/lsst/daf/butler/transfers/_context.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/python/lsst/daf/butler/transfers/_context.py b/python/lsst/daf/butler/transfers/_context.py index 7a9ad3a146..b880cec6e7 100644 --- a/python/lsst/daf/butler/transfers/_context.py +++ b/python/lsst/daf/butler/transfers/_context.py @@ -358,13 +358,20 @@ def _computeDatasetAssociations(self) -> dict[str, list[DatasetAssociation]]: collectionTypes = {CollectionType.TAGGED} if datasetType.isCalibration(): collectionTypes.add(CollectionType.CALIBRATION) - associationIter = self._butler.registry.queryDatasetAssociations( - datasetType, - collections=self._collections.keys(), + resolved_collections = self._butler._registry.queryCollections( + self._collections.keys(), + datasetType=datasetType, collectionTypes=collectionTypes, flattenChains=False, ) - for association in associationIter: - if association.ref.id in self._dataset_ids: - results[association.collection].append(association) + with self._butler._query() as query: + query = query.join_dataset_search(datasetType, resolved_collections) + result = query.general( + datasetType.dimensions, + dataset_fields={datasetType.name: {"dataset_id", "run", "collection", "timespan"}}, + ) + for association in DatasetAssociation.from_query_result(result, datasetType): + if association.ref.id in self._dataset_ids: + results[association.collection].append(association) + return results