def testPickle(self):
"""Test pickle support.
"""
storageClass = StorageClass("test_pickle")
datasetTypeName = "test"
dimensions = self.universe.extract(("instrument", "visit"))
# Un-pickling requires that storage class is registered with factory.
StorageClassFactory().registerStorageClass(storageClass)
datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
datasetTypeOut = pickle.loads(pickle.dumps(datasetType))
self.assertIsInstance(datasetTypeOut, DatasetType)
self.assertEqual(datasetType.name, datasetTypeOut.name)
self.assertEqual(datasetType.dimensions.names, datasetTypeOut.dimensions.names)
self.assertEqual(datasetType.storageClass, datasetTypeOut.storageClass)
self.assertIsNone(datasetTypeOut.parentStorageClass)
self.assertIs(datasetType.isCalibration(), datasetTypeOut.isCalibration())
self.assertFalse(datasetTypeOut.isCalibration())
datasetType = DatasetType(datasetTypeName, dimensions, storageClass, isCalibration=True)
datasetTypeOut = pickle.loads(pickle.dumps(datasetType))
self.assertIs(datasetType.isCalibration(), datasetTypeOut.isCalibration())
self.assertTrue(datasetTypeOut.isCalibration())
# And again with a composite
componentStorageClass = StorageClass("pickle_component")
StorageClassFactory().registerStorageClass(componentStorageClass)
componentDatasetType = DatasetType(DatasetType.nameWithComponent(datasetTypeName, "comp"),
dimensions, componentStorageClass,
parentStorageClass=storageClass)
datasetTypeOut = pickle.loads(pickle.dumps(componentDatasetType))
self.assertIsInstance(datasetTypeOut, DatasetType)
self.assertEqual(componentDatasetType.name, datasetTypeOut.name)
self.assertEqual(componentDatasetType.dimensions.names, datasetTypeOut.dimensions.names)
self.assertEqual(componentDatasetType.storageClass, datasetTypeOut.storageClass)
self.assertEqual(componentDatasetType.parentStorageClass, datasetTypeOut.parentStorageClass)
self.assertEqual(datasetTypeOut.parentStorageClass.name,
storageClass.name)
self.assertEqual(datasetTypeOut, componentDatasetType)
# Now with a string and not a real storage class to test that
# pickling doesn't force the StorageClass to be resolved
componentDatasetType = DatasetType(DatasetType.nameWithComponent(datasetTypeName, "comp"),
dimensions, "StrangeComponent",
parentStorageClass="UnknownParent")
datasetTypeOut = pickle.loads(pickle.dumps(componentDatasetType))
self.assertEqual(datasetTypeOut, componentDatasetType)
self.assertEqual(datasetTypeOut._parentStorageClassName,
componentDatasetType._parentStorageClassName)
# Now with a storage class that is created by the factory
factoryStorageClassClass = StorageClassFactory.makeNewStorageClass("ParentClass")
factoryComponentStorageClassClass = StorageClassFactory.makeNewStorageClass("ComponentClass")
componentDatasetType = DatasetType(DatasetType.nameWithComponent(datasetTypeName, "comp"),
dimensions, factoryComponentStorageClassClass(),
parentStorageClass=factoryStorageClassClass())
datasetTypeOut = pickle.loads(pickle.dumps(componentDatasetType))
self.assertEqual(datasetTypeOut, componentDatasetType)
self.assertEqual(datasetTypeOut._parentStorageClassName,
componentDatasetType._parentStorageClassName)
def register(self, datasetType: DatasetType) -> Tuple[DatasetRecordStorage, bool]:
# Docstring inherited from DatasetRecordStorageManager.
if datasetType.isComponent():
raise ValueError("Component dataset types can not be stored in registry."
f" Rejecting {datasetType.name}")
storage = self._byName.get(datasetType.name)
if storage is None:
dimensionsKey = self._dimensions.saveDimensionGraph(datasetType.dimensions)
tagTableName = makeTagTableName(datasetType, dimensionsKey)
calibTableName = (makeCalibTableName(datasetType, dimensionsKey)
if datasetType.isCalibration() else None)
row, inserted = self._db.sync(
self._static.dataset_type,
keys={"name": datasetType.name},
compared={
"dimensions_key": dimensionsKey,
"storage_class": datasetType.storageClass.name,
},
extra={
"tag_association_table": tagTableName,
"calibration_association_table": calibTableName,
},
returning=["id", "tag_association_table"],
)
assert row is not None
tags = self._db.ensureTableExists(
tagTableName,
makeTagTableSpec(datasetType, type(self._collections)),
)
if calibTableName is not None:
calibs = self._db.ensureTableExists(
calibTableName,
makeCalibTableSpec(datasetType, type(self._collections),
self._db.getTimespanRepresentation()),
)
else:
calibs = None
storage = ByDimensionsDatasetRecordStorage(db=self._db, datasetType=datasetType,
static=self._static, summaries=self._summaries,
tags=tags, calibs=calibs,
dataset_type_id=row["id"],
collections=self._collections)
self._byName[datasetType.name] = storage
self._byId[storage._dataset_type_id] = storage
else:
if datasetType != storage.datasetType:
raise ConflictingDefinitionError(f"Given dataset type {datasetType} is inconsistent "
f"with database definition {storage.datasetType}.")
inserted = False
return storage, inserted
def makeCalibTableName(datasetType: DatasetType, dimensionsKey: int) -> str:
"""Construct the name for a dynamic (DatasetType-dependent) tag + validity
range table used by the classes in this package.
Parameters
----------
datasetType : `DatasetType`
Dataset type to construct a name for. Multiple dataset types may
share the same table.
dimensionsKey : `int`
Integer key used to save ``datasetType.dimensions`` to the database.
Returns
-------
name : `str`
Name for the table.
"""
assert datasetType.isCalibration()
return f"dataset_calibs_{dimensionsKey:08d}"
