Make addDataset optionally recursive.

python/lsst/daf/butler/butler.py

@@ -237,26 +237,23 @@ def put(self, obj, datasetRefOrType, dataId=None, producer=None):
             if dataId is None:
                 raise ValueError("Must provide a dataId if first argument is not a DatasetRef")
             datasetType = self.registry.getDatasetType(datasetRefOrType)
-        ref = self.registry.addDataset(datasetType, dataId, run=self.run, producer=producer)
 
-        # Look up storage class to see if this is a composite
-        storageClass = datasetType.storageClass
+        isVirtualComposite = self.composites.doDisassembly(datasetType)
+
+        # Add Registry Dataset entry.  If not a virtual composite, add
+        # and attach components at the same time.
+        ref = self.registry.addDataset(datasetType, dataId, run=self.run, producer=producer,
+                                       recursive=not isVirtualComposite)
 
         # Check to see if this datasetType requires disassembly
-        if self.composites.doDisassembly(datasetType):
-            components = storageClass.assembler().disassemble(obj)
+        if isVirtualComposite:
+            components = datasetType.storageClass.assembler().disassemble(obj)
             for component, info in components.items():
                 compTypeName = datasetType.componentTypeName(component)
                 compRef = self.put(info.component, compTypeName, dataId, producer)
                 self.registry.attachComponent(component, ref, compRef)
         else:
             # This is an entity without a disassembler.
-            # If it is a composite we still need to register the components
-            for component in storageClass.components:
-                compTypeName = datasetType.componentTypeName(component)
-                compDatasetType = self.registry.getDatasetType(compTypeName)
-                compRef = self.registry.addDataset(compDatasetType, dataId, run=self.run, producer=producer)
-                self.registry.attachComponent(component, ref, compRef)
             self.datastore.put(obj, ref)
 
         return ref

python/lsst/daf/butler/core/registry.py

@@ -266,7 +266,7 @@ def getDatasetType(self, name):
 
     @abstractmethod
     @transactional
-    def addDataset(self, datasetType, dataId, run, producer=None):
+    def addDataset(self, datasetType, dataId, run, producer=None, recursive=False):
         """Adds a Dataset entry to the `Registry`
 
         This always adds a new Dataset; to associate an existing Dataset with
@@ -287,6 +287,9 @@ def addDataset(self, datasetType, dataId, run, producer=None):
             Unit of work that produced the Dataset.  May be ``None`` to store
             no provenance information, but if present the `Quantum` must
             already have been added to the Registry.
+        recursive : `bool`
+            If True, recursively add Dataset and attach entries for component
+            Datasets as well.
 
         Returns
         -------

python/lsst/daf/butler/registries/sqlRegistry.py

@@ -340,7 +340,7 @@ def getDatasetType(self, name):
         return datasetType
 
     @transactional
-    def addDataset(self, datasetType, dataId, run, producer=None):
+    def addDataset(self, datasetType, dataId, run, producer=None, recursive=False):
         """Adds a Dataset entry to the `Registry`
 
         This always adds a new Dataset; to associate an existing Dataset with
@@ -361,6 +361,9 @@ def addDataset(self, datasetType, dataId, run, producer=None):
             Unit of work that produced the Dataset.  May be ``None`` to store
             no provenance information, but if present the `Quantum` must
             already have been added to the SqlRegistry.
+        recursive : `bool`
+            If True, recursively add Dataset and attach entries for component
+            Datasets as well.
 
         Returns
         -------
@@ -396,6 +399,14 @@ def addDataset(self, datasetType, dataId, run, producer=None):
                                 run=run)
         # A dataset is always associated with its Run collection
         self.associate(run.collection, [datasetRef, ])
+
+        if recursive:
+            for component in datasetType.storageClass.components:
+                compTypeName = datasetType.componentTypeName(component)
+                compDatasetType = self.getDatasetType(compTypeName)
+                compRef = self.addDataset(compDatasetType, dataId, run=run, producer=producer,
+                                          recursive=True, transactional=False)
+                self.attachComponent(component, datasetRef, compRef)
         return datasetRef
 
     def getDataset(self, id):