Merge pull request #41 from lsst/tickets/DM-14520

DM-14520: re-enable (and reimplement) Dataset-DataUnit foreign key constraints
lsst · May 19, 2018 · 2316b8d · 2316b8d
2 parents 39b8556 + 388f55c
commit 2316b8d
Show file tree

Hide file tree

Showing 5 changed files with 92 additions and 14 deletions.
diff --git a/config/schema.yaml b/config/schema.yaml
@@ -10,7 +10,9 @@ schema:
       -
         name: camera
         type: string
-        #foreign_key: Camera.camera
+        foreignKey:
+          src: camera
+          tgt: Camera.camera
       tables:
         Camera:
           doc: >
@@ -47,7 +49,13 @@ schema:
       -
         name: physical_filter
         type: string
-        #foreign_key: PhysicalFilter.physical_filter
+        foreignKey:
+          src:
+            - camera
+            - physical_filter
+          tgt:
+            - PhysicalFilter.camera
+            - PhysicalFilter.physical_filter
       tables:
         PhysicalFilter:
           doc: >
@@ -88,7 +96,13 @@ schema:
       -
         name: sensor
         type: string
-        #foreign_key: Sensor.sensor
+        foreignKey:
+          src:
+            - camera
+            - sensor
+          tgt:
+            - Sensor.camera
+            - Sensor.sensor
       tables:
         Sensor:
           doc: >
@@ -148,7 +162,13 @@ schema:
       -
         name: visit
         type: int
-        #foreign_key: Visit.visit
+        foreignKey:
+          src:
+            - camera
+            - visit
+          tgt:
+            - Visit.camera
+            - Visit.visit
       tables:
         Visit:
           doc: >
@@ -349,7 +369,13 @@ schema:
       -
         name: exposure
         type: int
-        #foreign_key: Exposure.exposure
+        foreignKey:
+          src:
+            - camera
+            - exposure
+          tgt:
+            - Exposure.camera
+            - Exposure.exposure
       tables:
         Exposure:
           doc: >
@@ -499,7 +525,9 @@ schema:
       -
         name: skymap
         type: string
-        #foreign_key: SkyMap.skymap
+        foreignKey:
+          src: skymap
+          tgt: SkyMap.skymap
       tables:
         SkyMap:
           doc: >
@@ -533,7 +561,13 @@ schema:
       -
         name: tract
         type: int
-        #foreign_key: Tract.tract
+        foreignKey:
+          src:
+            - skymap
+            - tract
+          tgt:
+            - Tract.skymap
+            - Tract.tract
       tables:
         Tract:
           doc: >
@@ -578,7 +612,15 @@ schema:
       -
         name: patch
         type: int
-        #foreign_key: Patch.patch
+        foreignKey:
+          src:
+            - skymap
+            - tract
+            - patch
+          tgt:
+            - Patch.skymap
+            - Patch.tract
+            - Patch.patch
       tables:
         Patch:
           doc: >

diff --git a/python/lsst/daf/butler/core/dataUnit.py b/python/lsst/daf/butler/core/dataUnit.py
@@ -160,6 +160,7 @@ def __init__(self):
         self._dataUnitNames = None
         self._dataUnits = {}
         self.links = {}
+        self.constraints = []
 
     @classmethod
     def fromConfig(cls, config, builder=None):
@@ -253,7 +254,12 @@ def _initDataUnits(self, config, builder):
                     link = tuple((linkDescription['name'] for linkDescription in dataUnitDescription['link']))
                     # Link columns that will become part of the Datasets table
                     for linkDescription in dataUnitDescription['link']:
-                        self.links[linkDescription['name']] = builder.makeColumn(linkDescription)
+                        linkColumnDesc = linkDescription.copy()
+                        linkConstraintDesc = linkColumnDesc.pop("foreignKey", None)
+                        linkName = linkDescription['name']
+                        self.links[linkName] = builder.makeColumn(linkColumnDesc)
+                        if linkConstraintDesc is not None:
+                            self.constraints.append(builder.makeForeignKeyConstraint(linkConstraintDesc))
                 if 'tables' in dataUnitDescription:
                     for tableName, tableDescription in dataUnitDescription['tables'].items():
                         if tableName == dataUnitName:

diff --git a/python/lsst/daf/butler/core/schema.py b/python/lsst/daf/butler/core/schema.py
@@ -22,7 +22,7 @@
 from .utils import iterable
 from .config import ConfigSubset
 from sqlalchemy import Column, String, Integer, Boolean, LargeBinary, DateTime,\
-    Float, ForeignKey, ForeignKeyConstraint, Table, MetaData
+    Float, ForeignKeyConstraint, Table, MetaData
 from .dataUnit import DataUnitRegistry
 
 metadata = None  # Needed to make disabled test_hsc not fail on import
@@ -64,6 +64,8 @@ def buildFromConfig(self, config):
         datasetTable = self.builder.metadata.tables['Dataset']
         for linkColumn in self.dataUnits.links.values():
             datasetTable.append_column(linkColumn)
+        for linkConstraint in self.dataUnits.constraints:
+            datasetTable.append_constraint(linkConstraint)
         self.metadata = self.builder.metadata
 
 
@@ -163,7 +165,6 @@ def makeColumn(self, columnDescription):
             May contain:
             - nullable, entry can be null
             - primary_key, mark this column as primary key
-            - foreign_key, link to other table
             - doc, docstring
 
         Returns
@@ -180,9 +181,6 @@ def makeColumn(self, columnDescription):
         # required
         columnName = description.pop("name")
         args = (columnName, self.VALID_COLUMN_TYPES[description.pop("type")])
-        # foreign_key is special
-        if "foreign_key" in description:
-            args += (ForeignKey(description.pop("foreign_key")), )
         # additional optional arguments can be passed through directly
         kwargs = {}
         for opt in ("nullable", "primary_key", "doc"):

diff --git a/tests/test_butler.py b/tests/test_butler.py
@@ -84,6 +84,12 @@ def testBasicPutGet(self):
         storageClass = self.storageClassFactory.getStorageClass("StructuredData")
         self.registerDatasetTypes(datasetTypeName, dataUnits, storageClass, butler.registry)
 
+        # Add needed DataUnits
+        butler.registry.addDataUnitEntry("Camera", {"camera": "DummyCam"})
+        butler.registry.addDataUnitEntry("PhysicalFilter", {"camera": "DummyCam", "physical_filter": "d-r"})
+        butler.registry.addDataUnitEntry("Visit", {"camera": "DummyCam", "visit": 42,
+                                                   "physical_filter": "d-r"})
+
         # Create and store a dataset
         metric = makeExampleMetrics()
         dataId = {"camera": "DummyCam", "visit": 42}
@@ -108,6 +114,13 @@ def testCompositePutGet(self):
         storageClass = self.storageClassFactory.getStorageClass("StructuredComposite")
         self.registerDatasetTypes(datasetTypeName, dataUnits, storageClass, butler.registry)
 
+        # Add needed DataUnits
+        butler.registry.addDataUnitEntry("Camera", {"camera": "DummyCamComp"})
+        butler.registry.addDataUnitEntry("PhysicalFilter", {"camera": "DummyCamComp",
+                                                            "physical_filter": "d-r"})
+        butler.registry.addDataUnitEntry("Visit", {"camera": "DummyCamComp", "visit": 423,
+                                                   "physical_filter": "d-r"})
+
         # Create and store a dataset
         metric = makeExampleMetrics()
         dataId = {"camera": "DummyCamComp", "visit": 423}

diff --git a/tests/test_sqlRegistry.py b/tests/test_sqlRegistry.py
@@ -82,6 +82,7 @@ def testDataset(self):
         storageClass = StorageClass("testDataset")
         datasetType = DatasetType(name="testtype", dataUnits=("Camera",), storageClass=storageClass)
         registry.registerDatasetType(datasetType)
+        registry.addDataUnitEntry("Camera", {"camera": "DummyCam"})
         ref = registry.addDataset(datasetType, dataId={"camera": "DummyCam"}, run=run)
         outRef = registry.getDataset(ref.id)
         self.assertEqual(ref, outRef)
@@ -95,6 +96,7 @@ def testComponents(self):
         registry.registerDatasetType(parentDatasetType)
         registry.registerDatasetType(childDatasetType1)
         registry.registerDatasetType(childDatasetType2)
+        registry.addDataUnitEntry("Camera", {"camera": "DummyCam"})
         run = registry.makeRun(collection="test")
         parent = registry.addDataset(parentDatasetType, dataId={"camera": "DummyCam"}, run=run)
         children = {"child1": registry.addDataset(childDatasetType1, dataId={"camera": "DummyCam"}, run=run),
@@ -153,6 +155,7 @@ def testExecution(self):
 
     def testQuantum(self):
         registry = Registry.fromConfig(self.butlerConfig)
+        registry.addDataUnitEntry("Camera", {"camera": "DummyCam"})
         run = registry.makeRun(collection="test")
         storageClass = StorageClass("testQuantum")
         # Make two predicted inputs
@@ -184,6 +187,7 @@ def testStorageInfo(self):
         storageClass = StorageClass("testStorageInfo")
         datasetType = DatasetType(name="test", dataUnits=("Camera",), storageClass=storageClass)
         registry.registerDatasetType(datasetType)
+        registry.addDataUnitEntry("Camera", {"camera": "DummyCam"})
         run = registry.makeRun(collection="test")
         ref = registry.addDataset(datasetType, dataId={"camera": "DummyCam"}, run=run)
         datastoreName = "dummystore"
@@ -206,6 +210,7 @@ def testAssembler(self):
         storageClass = StorageClass("testAssembler")
         datasetType = DatasetType(name="test", dataUnits=("Camera",), storageClass=storageClass)
         registry.registerDatasetType(datasetType)
+        registry.addDataUnitEntry("Camera", {"camera": "DummyCam"})
         run = registry.makeRun(collection="test")
         ref = registry.addDataset(datasetType, dataId={"camera": "DummyCam"}, run=run)
         self.assertIsNone(ref.assembler)
@@ -219,6 +224,14 @@ def testFind(self):
         storageClass = StorageClass("testFind")
         datasetType = DatasetType(name="dummytype", dataUnits=("Camera", "Visit"), storageClass=storageClass)
         registry.registerDatasetType(datasetType)
+        registry.addDataUnitEntry("Camera", {"camera": "DummyCam"})
+        registry.addDataUnitEntry("Camera", {"camera": "MyCam"})
+        registry.addDataUnitEntry("PhysicalFilter", {"camera": "DummyCam", "physical_filter": "d-r"})
+        registry.addDataUnitEntry("PhysicalFilter", {"camera": "MyCam", "physical_filter": "m-r"})
+        registry.addDataUnitEntry("Visit", {"camera": "DummyCam", "visit": 0, "physical_filter": "d-r"})
+        registry.addDataUnitEntry("Visit", {"camera": "DummyCam", "visit": 1, "physical_filter": "d-r"})
+        registry.addDataUnitEntry("Visit", {"camera": "DummyCam", "visit": 2, "physical_filter": "d-r"})
+        registry.addDataUnitEntry("Visit", {"camera": "MyCam", "visit": 2, "physical_filter": "m-r"})
         collection = "test"
         dataId = {"camera": "DummyCam", "visit": 0}
         run = registry.makeRun(collection=collection)
@@ -251,8 +264,14 @@ def testCollections(self):
         storageClass = StorageClass("testCollections")
         datasetType = DatasetType(name="dummytype", dataUnits=("Camera", "Visit"), storageClass=storageClass)
         registry.registerDatasetType(datasetType)
+        registry.addDataUnitEntry("Camera", {"camera": "DummyCam"})
+        registry.addDataUnitEntry("PhysicalFilter", {"camera": "DummyCam", "physical_filter": "d-r"})
+        registry.addDataUnitEntry("Visit", {"camera": "DummyCam", "visit": 0, "physical_filter": "d-r"})
+        registry.addDataUnitEntry("Visit", {"camera": "DummyCam", "visit": 1, "physical_filter": "d-r"})
         collection = "ingest"
         run = registry.makeRun(collection=collection)
+        # TODO: Dataset.physical_filter should be populated as well here
+        # from the Visit DataUnit values.
         dataId1 = {"camera": "DummyCam", "visit": 0}
         inputRef1 = registry.addDataset(datasetType, dataId=dataId1, run=run)
         dataId2 = {"camera": "DummyCam", "visit": 1}