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}"