#!/usr/bin/env python import copy import types import re from dm.common.objects.dmObject import DmObject from dm.common.exceptions.dmException import DmException from dm.common.mongodb.api.dmMongoDbApi import DmMongoDbApi from dm.common.mongodb.impl.datasetCollection import DatasetCollection from dm.common.mongodb.api.fileMongoDbApi import FileMongoDbApi from dm.common.mongodb.api.dmMongoDbApi import DmMongoDbApi from dm.common.objects.datasetMetadata import DatasetMetadata class DatasetMongoDbApi(DmMongoDbApi): SYSTEM_KEY_LIST = ['_id', 'datasetName', 'experimentName'] USE_DATASET_NAME_KEY = '_useDatasetName' USE_EXPERIMENT_NAME_KEY = '_useDatasetName' def __init__(self): DmMongoDbApi.__init__(self) self.datasetCollection = DatasetCollection(self.dbClient) self.fileMongoDbApi = FileMongoDbApi() @DmMongoDbApi.executeDbCall def addExperimentDataset(self, datasetInfo, **kwargs): datasetInfo2 = self.getMongoDict(datasetInfo) dbDatasetMetadata = self.datasetCollection.addByUniqueKeys(datasetInfo2) return self.toDmObject(dbDatasetMetadata, DatasetMetadata) @DmMongoDbApi.executeDbCall def getDatasets(self, queryDict={}, returnFieldDict=DatasetCollection.ALL_FIELDS_DICT, **kwargs): queryDict2 = self.getMongoDict(queryDict) return self.listToDmObjects(self.datasetCollection.findByQueryDict(queryDict2, returnFieldDict), DatasetMetadata) @DmMongoDbApi.executeDbCall def getDatasetById(self, id, **kwargs): dbDatasetMetadata = self.datasetCollection.findById(id) return self.toDmObject(dbDatasetMetadata, DatasetMetadata) @DmMongoDbApi.executeDbCall def getExperimentDataset(self, experimentName, datasetName, **kwargs): queryDict = { 'datasetName' : datasetName, 'experimentName' : experimentName } dbDatasetMetadata = self.datasetCollection.findByUniqueKeys(queryDict) return self.toDmObject(dbDatasetMetadata, DatasetMetadata) @DmMongoDbApi.executeDbCall def getExperimentDatasets(self, experimentName, queryDict={}, returnFieldDict=DatasetCollection.ALL_FIELDS_DICT, **kwargs): queryDict2 = copy.copy(queryDict) queryDict2['experimentName'] = experimentName return self.getDatasets(queryDict2, returnFieldDict) @DmMongoDbApi.executeDbCall def updateDatasetById(self, datasetInfo, **kwargs): datasetInfo2 = self.getMongoDict(datasetInfo) dbDatasetMetadata = self.datasetCollection.updateById(datasetInfo2) return self.toDmObject(dbDatasetMetadata, DatasetMetadata) @DmMongoDbApi.executeDbCall def updateExperimentDataset(self, datasetInfo, **kwargs): datasetInfo2 = self.getMongoDict(datasetInfo) dbDatasetMetadata = self.datasetCollection.updateByUniqueKeys(datasetInfo2) return self.toDmObject(dbDatasetMetadata, DatasetMetadata) @DmMongoDbApi.executeDbCall def updateOrAddExperimentDataset(self, datasetInfo, **kwargs): datasetInfo2 = self.getMongoDict(datasetInfo) dbDatasetMetadata = self.datasetCollection.updateOrAddByUniqueKeys(datasetInfo2) return self.toDmObject(dbDatasetMetadata, DatasetMetadata) @DmMongoDbApi.executeDbCall def getExperimentDatasetFiles(self, experimentName, datasetName, returnFieldDict=DatasetCollection.ALL_FIELDS_DICT, **kwargs): self.getLogger().debug('Looking for dataset %s (experiment %s)' % (datasetName, experimentName)) datasetDict = self.getExperimentDataset(experimentName, datasetName).getDictRep(DmObject.ALL_KEYS) del datasetDict['id'] useDatasetName = kwargs.get(self.USE_DATASET_NAME_KEY, False) if not useDatasetName: del datasetDict['datasetName'] useExperimentName = kwargs.get(self.USE_EXPERIMENT_NAME_KEY, False) if not useExperimentName: del datasetDict['experimentName'] self.getLogger().debug('Converting dataset dict to regex patterns') ignoreCase = kwargs.get(self.REGEX_IGNORE_CASE_KEY, True) queryDict = self.convertStringsToRegex(datasetDict, ignoreCase) return self.fileMongoDbApi.getExperimentFiles(experimentName, queryDict, returnFieldDict) ####################################################################### # Testing. if __name__ == '__main__': api = DatasetMongoDbApi() import time t = long(time.time()) datasetName = 'd-%s' % t experimentName = 'exp-01' datasetInfo = {'datasetName' : datasetName, 'intKey' : 1, 'doubleKey' : 2.0, 'stringKey' : 'myString' , 'experimentName' : experimentName} dataset = api.updateOrAddExperimentDataset(datasetInfo) print '\nADDED DATASET\n', dataset datasets = api.getDatasets() print '\nDATASETS: \n', api.getDatasets() for dataset in datasets: print dataset.getDictRep() #import re #regex = re.compile("file0.*", re.IGNORECASE) datasetName = 'dataset1' datasetInfo = {'datasetName' : datasetName, 'experimentName' : experimentName, 'fileName' : 'file0.*'} dataset = api.updateOrAddExperimentDataset(datasetInfo) print '\nUPDATED DATASET\n', dataset print '\nDATASET FILES\n', api.getExperimentDatasetFiles(experimentName, datasetName)