Forked from
DM / dm-docs
261 commits behind, 293 commits ahead of the upstream repository.
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
dmMongoCollection.py 8.70 KiB
#!/usr/bin/env python
from bson.objectid import ObjectId
from dm.common.utility.loggingManager import LoggingManager
from dm.common.exceptions.invalidArgument import InvalidArgument
from dm.common.exceptions.objectAlreadyExists import ObjectAlreadyExists
from dm.common.exceptions.objectNotFound import ObjectNotFound
from dm.common.exceptions.dbError import DbError
class DmMongoCollection(object):
"""Base collection in mongo db."""
ALL_FIELDS_DICT = {'__return_only_id__' : False}
UNIQUE_KEYS_LIST = []
def __init__(self, collectionName, dbClient):
self.logger = LoggingManager.getInstance().getLogger(self.__class__.__name__)
self.dbClient = dbClient
# Collection consists of items: [item]
self.collectionName = collectionName
# items => item
self.itemName = collectionName[:-1]
# item => Item
self.capitalizedItemName = self.itemName.capitalize()
@classmethod
def getUniqueKeys(cls):
return cls.UNIQUE_KEYS_LIST
def findByKey(self, key, value):
dbObjectDict = self.dbClient.findOne(self.collectionName, {key : value})
if dbObjectDict is None:
raise ObjectNotFound('%s with %s=%s not found.' % (self.capitalizedItemName, key, value))
return dbObjectDict
def findByKeys(self, keyList, queryDict):
for key in keyList:
if not queryDict.has_key(key):
raise InvalidArgument('%s query dictionary does not specify key %s.' % (self.capitalizedItemName, key))
dbObjectDict = self.dbClient.findOne(self.collectionName, queryDict)
if dbObjectDict is None:
raise ObjectNotFound('%s with properties %s not found.' % (self.capitalizedItemName, queryDict))
return dbObjectDict
def findByUniqueKeys(self, queryDict):
return self.findByKeys(self.UNIQUE_KEYS_LIST, queryDict)
def findById(self, id):
return self.findByKey('_id', ObjectId(id))
def findByName(self, name):
return self.findByKey('name', name)
def findByQueryDict(self, queryDict, returnFieldDict=ALL_FIELDS_DICT):
return self.dbClient.findAsList(self.collectionName, queryDict, returnFieldDict)
def listByKey(self, key):
return self.dbClient.findAsList(self.collectionName, {}, {key : True})
def listById(self):
return self.dbClient.findAsList(self.collectionName, {}, {})
def listByName(self):
return self.findByName('name')
def listByKeys(self, keyList):
returnFieldDict = {}
for key in keyList:
returnFieldDict[key] = True
return self.dbClient.findAsList(self.collectionName, {}, returnFieldDict)
def listAll(self):
return self.dbClient.findAsList(self.collectionName, {}, DmMongoCollection.ALL_FIELDS_DICT)
def __addDbObject(self, objectDict):
try:
self.dbClient.insert(self.collectionName, objectDict)
except Exception, ex:
self.logger.error('Cannot add %s with %s set to %s: %s' % (self.itemName, key, value, ex))
raise DbError(exception=ex)
return objectDict
def addByKey(self, key, objectDict):
value = objectDict.get(key)
if value is None:
raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
dbObjectDict = self.dbClient.findOne(self.collectionName, {key : value})
if dbObjectDict is not None:
raise ObjectAlreadyExists('%s with %s set to %s already exists.' % (self.capitalizedItemName, key, value))
return self.__addDbObject(objectDict)
def addByName(self, objectDict):
return self.addByKey('name', objectDict)
def addByKeys(self, keyList, objectDict):
queryDict = {}
for key in keyList:
value = objectDict.get(key)
if value is None:
raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
queryDict[key] = value
dbObjectDict = self.dbClient.findOne(self.collectionName, queryDict)
if dbObjectDict is not None:
raise ObjectAlreadyExists('%s with properties %s already exists.' % (self.capitalizedItemName, queryDict))
return self.__addDbObject(objectDict)
def addByUniqueKeys(self, objectDict):
return self.addByKeys(self.UNIQUE_KEYS_LIST, objectDict)
def __updateDbObject(self, dbObjectDict, objectDict):
try:
id = dbObjectDict.get('_id')
dbObjectDict.update(objectDict)
self.dbClient.update(self.collectionName,
{'_id' : id},
{'$set' : dbObjectDict})
except Exception, ex:
self.logger.error('Cannot update %s %s with %s: %s' % (self.itemName, dbObjectDict, objectDict, ex))
raise DbError(exception=ex)
return dbObjectDict
def updateByKey(self, key, objectDict):
value = objectDict.get(key)
if value is None:
raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
dbObjectDict = self.dbClient.findOne(self.collectionName, {key : value})
if dbObjectDict is None:
raise ObjectNotFound('%s with %s set to %s not found.' % (self.capitalizedItemName, key, value))
return self.__updateDbObject(dbObjectDict, objectDict)
def updateByKeys(self, keyList, objectDict):
queryDict = {}
for key in keyList:
value = objectDict.get(key)
if value is None:
raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
queryDict[key] = value
dbObjectDict = self.dbClient.findOne(self.collectionName, queryDict)
if dbObjectDict is None:
raise ObjectNotFound('%s with properties %s not found.' % (self.capitalizedItemName, queryDict))
return self.__updateDbObject(dbObjectDict, objectDict)
def updateByUniqueKeys(self, objectDict):
return self.updateByKeys(self.UNIQUE_KEYS_LIST, objectDict)
def updateById(self, objectDict):
# Convert 'id' to '_id' key if needed, and wrap it with ObjectId()
if not objectDict.has_key('_id') and objectDict.has_key('id'):
objectDict['_id'] = ObjectId(objectDict['id'])
del objectDict['id']
return self.updateByKey('_id', objectDict)
def updateByName(self, objectDict):
return self.updateByKey('name', objectDict)
def updateOrAddByKey(self, key, objectDict):
value = objectDict.get(key)
if value is None:
raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
dbObjectDict = self.dbClient.findOne(self.collectionName, {key : value})
if dbObjectDict is None:
return self.__addDbObject(objectDict)
else:
return self.__updateDbObject(dbObjectDict, objectDict)
def updateOrAddByKeys(self, keyList, objectDict):
queryDict = {}
for key in keyList:
value = objectDict.get(key)
if value is None:
raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
queryDict[key] = value
dbObjectDict = self.dbClient.findOne(self.collectionName, queryDict)
if dbObjectDict is None:
return self.__addDbObject(objectDict)
else:
return self.__updateDbObject(dbObjectDict, objectDict)
def updateOrAddByUniqueKeys(self, objectDict):
return self.updateOrAddByKeys(self.UNIQUE_KEYS_LIST, objectDict)
def updateOrAddByName(self, objectDict):
return self.updateOrAddByKey('name', objectDict)
#######################################################################
# Testing
if __name__ == '__main__':
from dmMongoClient import DmMongoClient
from bson.objectid import ObjectId
mongo = DmMongoClient('dm')
fileCollectionImpl = DmMongoCollection('files', mongo)
#objectDict = {'name' : 'xyz-001', 'experiment' : 'myexp-D', 'update' : 'sv2', 'locationList' : '[/opt/xyz, /data/xyz]'}
#print fileCollectionImpl.updateOrAddByKey('name', objectDict)
#f = fileCollectionImpl.findByKey('name', 'xyz-001')
#print f
#print type(f['_id'])
#print fileCollectionImpl.findById('556de0059e058b0ef4c4413b')
#print fileCollectionImpl.findByQueryDict({'experiment' : 'exp-001'}, {'locationList' : 1})
#print 'LIST BY ID\n', fileCollectionImpl.listById()
#print 'LIST BY NAME\n', fileCollectionImpl.listByKey('name')
print 'LIST ALL\n'
for f in fileCollectionImpl.listAll():
print 'FILE: %s\n' % f
#print fileCollectionImpl.updateByKeys(['name', 'experiment'], objectDict)