"""Event Man(ager) database backend Classes and functions used to manage events and attendees database. Copyright 2015 Davide Alberani RaspiBO Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ import re import pymongo from bson.objectid import ObjectId re_objectid = re.compile(r'[0-9a-f]{24}') class EventManDB(object): """MongoDB connector.""" db = None connection = None # map operations on lists of items. _operations = { 'update': '$set', 'append': '$push', 'appendUnique': '$addToSet', 'delete': '$pull' } def __init__(self, url=None, dbName='eventman'): """Initialize the instance, connecting to the database. :param url: URL of the database :type url: str (or None to connect to localhost) """ self._url = url self._dbName = dbName self.connect(url) def connect(self, url=None, dbName=None): """Connect to the database. :param url: URL of the database :type url: str (or None to connect to localhost) :return: the database we're connected to :rtype: :class:`~pymongo.database.Database` """ if self.db is not None: return self.db if url: self._url = url if dbName: self._dbName = dbName self.connection = pymongo.MongoClient(self._url) self.db = self.connection[self._dbName] return self.db def convert_obj(self, obj): """Convert an object in a format suitable to be stored in MongoDB. :param obj: object to convert :return: object that can be stored in MongoDB. """ if obj is None: return None try: return ObjectId(obj) except: pass try: return int(obj) except: pass return obj def convert(self, seq): """Convert an object in a format suitable to be stored in MongoDB, descending lists, tuples and dictionaries (a copy is returned). :param seq: sequence or object to convert :return: object that can be stored in MongoDB. """ if isinstance(seq, dict): d = {} for key, item in seq.iteritems(): d[key] = self.convert(item) return d if isinstance(seq, (list, tuple)): return [self.convert(x) for x in seq] return self.convert_obj(seq) def get(self, collection, _id): """Get a single document with the specified `_id`. :param collection: search the document in this collection :type collection: str :param _id: unique ID of the document :type _id: str or :class:`~bson.objectid.ObjectId` :return: the document with the given `_id` :rtype: dict """ results = self.query(collection, self.convert({'_id': _id})) return results and results[0] or {} def query(self, collection, query=None): """Get multiple documents matching a query. :param collection: search for documents in this collection :type collection: str :param query: search for documents with those attributes :type query: dict or None :return: list of matching documents :rtype: list """ db = self.connect() query = self.convert(query or {}) return list(db[collection].find(query)) def add(self, collection, data): """Insert a new document. :param collection: insert the document in this collection :type collection: str :param data: the document to store :type data: dict :return: the document, as created in the database :rtype: dict """ db = self.connect() data = self.convert(data) _id = db[collection].insert(data) return self.get(collection, _id) def insertOne(self, collection, data): """Insert a document, avoiding duplicates. :param collection: update a document in this collection :type collection: str :param data: the document information to store :type data: dict :return: True if the document was already present :rtype: bool """ db = self.connect() data = self.convert(data) ret = db[collection].update(data, {'$set': data}, upsert=True) return ret['updatedExisting'] def _buildSearchPattern(self, data, searchBy): """Return an OR condition.""" _or = [] for searchPattern in searchBy: try: _or.append(dict([(k, data[k]) for k in searchPattern if k in data])) except KeyError: continue return _or def update(self, collection, _id_or_query, data, operation='update', updateList=None, create=True): """Update an existing document or create it, if requested. _id_or_query can be an ID, a dict representing a query or a list of tuples. In the latter case, the tuples are put in OR; a tuple match if all of its items from 'data' are contained in the document. :param collection: update a document in this collection :type collection: str :param _id_or_query: ID of the document to be updated, or a query or a list of attributes in the data that must match :type _id_or_query: str or :class:`~bson.objectid.ObjectId` or iterable :param data: the updated information to store :type data: dict :param operation: operation used to update the document or a portion of it, like a list (update, append, appendUnique, delete) :type operation: str :param updateList: if set, it's considered the name of a list (the first matching element will be updated) :type updateList: str :param create: if True, the document is created if no document matches :type create: bool :return: a boolean (True if an existing document was updated) and the document after the update :rtype: tuple of (bool, dict) """ db = self.connect() data = self.convert(data or {}) _id_or_query = self.convert(_id_or_query) if isinstance(_id_or_query, (list, tuple)): _id_or_query = {'$or': self._buildSearchPattern(data, _id_or_query)} elif not isinstance(_id_or_query, dict): _id_or_query = {'_id': _id_or_query} if '_id' in data: del data['_id'] operator = self._operations.get(operation) if updateList: newData = {} for key, value in data.iteritems(): newData['%s.$.%s' % (updateList, key)] = value data = newData res = db[collection].find_and_modify(query=_id_or_query, update={operator: data}, full_response=True, new=True, upsert=create) lastErrorObject = res.get('lastErrorObject') or {} return lastErrorObject.get('updatedExisting', False), res.get('value') or {} def delete(self, collection, _id_or_query=None, force=False): """Remove one or more documents from a collection. :param collection: search the documents in this collection :type collection: str :param _id_or_query: unique ID of the document or query to match multiple documents :type _id_or_query: str or :class:`~bson.objectid.ObjectId` or dict :param force: force the deletion of all documents, when `_id_or_query` is empty :type force: bool """ if not _id_or_query and not force: return db = self.connect() if not isinstance(_id_or_query, dict): _id_or_query = {'_id': _id_or_query} _id_or_query = self.convert(_id_or_query) db[collection].remove(_id_or_query)