311 строки
11 КиБ
Python
311 строки
11 КиБ
Python
# -*- coding: utf-8 -*-
|
|
"""Monco: a MongoDB database backend
|
|
|
|
Classes and functions used to issue queries to a MongoDB database.
|
|
|
|
Copyright 2016-2017 Davide Alberani <da@erlug.linux.it>
|
|
RaspiBO <info@raspibo.org>
|
|
|
|
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}')
|
|
|
|
_force_conversion = {
|
|
'_id': ObjectId,
|
|
'seq_hex': str,
|
|
'tickets.seq_hex': str
|
|
}
|
|
|
|
|
|
def convert_obj(obj):
|
|
"""Convert an object in a format suitable to be stored in MongoDB.
|
|
|
|
:param obj: object to convert
|
|
|
|
:returns: object that can be stored in MongoDB.
|
|
"""
|
|
if obj is None:
|
|
return None
|
|
if isinstance(obj, bool):
|
|
return obj
|
|
try:
|
|
if re_objectid.match(obj):
|
|
return ObjectId(obj)
|
|
except:
|
|
pass
|
|
return obj
|
|
|
|
|
|
def convert(seq):
|
|
"""Convert an object to a format suitable to be stored in MongoDB,
|
|
descending lists, tuples and dictionaries (a copy is returned).
|
|
|
|
:param seq: sequence or object to convert
|
|
|
|
:returns: object that can be stored in MongoDB.
|
|
"""
|
|
if isinstance(seq, dict):
|
|
d = {}
|
|
for key, item in seq.items():
|
|
if key in _force_conversion:
|
|
try:
|
|
d[key] = _force_conversion[key](item)
|
|
except:
|
|
d[key] = item
|
|
else:
|
|
d[key] = convert(item)
|
|
return d
|
|
if isinstance(seq, (list, tuple)):
|
|
return [convert(x) for x in seq]
|
|
return convert_obj(seq)
|
|
|
|
|
|
class MoncoError(Exception):
|
|
"""Base class for Monco exceptions."""
|
|
pass
|
|
|
|
|
|
class MoncoConnectionError(MoncoError):
|
|
"""Monco exceptions raise when a connection problem occurs."""
|
|
pass
|
|
|
|
|
|
class Monco(object):
|
|
"""MongoDB connector."""
|
|
db = None
|
|
connection = None
|
|
|
|
# map operations on lists of items.
|
|
_operations = {
|
|
'update': '$set',
|
|
'append': '$push',
|
|
'appendUnique': '$addToSet',
|
|
'delete': '$pull',
|
|
'increment': '$inc'
|
|
}
|
|
|
|
def __init__(self, dbName, url=None):
|
|
"""Initialize the instance, connecting to the database.
|
|
|
|
:param dbName: name of the database
|
|
:type dbName: str (or None to use the dbName passed at initialization)
|
|
:param url: URL of the database
|
|
:type url: str (or None to connect to localhost)
|
|
"""
|
|
self._url = url
|
|
self._dbName = dbName
|
|
self.connect()
|
|
|
|
def connect(self, dbName=None, url=None):
|
|
"""Connect to the database.
|
|
|
|
:param dbName: name of the database
|
|
:type dbName: str (or None to use the dbName passed at initialization)
|
|
:param url: URL of the database
|
|
:type url: str (or None to connect to localhost)
|
|
|
|
:returns: 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
|
|
if not self._dbName:
|
|
raise MoncoConnectionError('no database name specified')
|
|
self.connection = pymongo.MongoClient(self._url)
|
|
self.db = self.connection[self._dbName]
|
|
return self.db
|
|
|
|
def getOne(self, collection, query=None):
|
|
"""Get a single document with the specified `query`.
|
|
|
|
:param collection: search the document in this collection
|
|
:type collection: str
|
|
:param query: query to filter the documents
|
|
:type query: dict or None
|
|
|
|
:returns: the first document matching the query
|
|
:rtype: dict
|
|
"""
|
|
results = self.query(collection, convert(query))
|
|
return results and results[0] or {}
|
|
|
|
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`
|
|
|
|
:returns: the document with the given `_id`
|
|
:rtype: dict
|
|
"""
|
|
return self.getOne(collection, {'_id': _id})
|
|
|
|
def query(self, collection, query=None, condition='or'):
|
|
"""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, list or None
|
|
|
|
:returns: list of matching documents
|
|
:rtype: list
|
|
"""
|
|
db = self.connect()
|
|
query = convert(query or {})
|
|
if isinstance(query, (list, tuple)):
|
|
query = {'$%s' % condition: query}
|
|
return list(db[collection].find(query))
|
|
|
|
def add(self, collection, data, _id=None):
|
|
"""Insert a new document.
|
|
|
|
:param collection: insert the document in this collection
|
|
:type collection: str
|
|
:param data: the document to store
|
|
:type data: dict
|
|
:param _id: the _id of the document to store; if None, it's generated
|
|
:type _id: object
|
|
|
|
:returns: the document, as created in the database
|
|
:rtype: dict
|
|
"""
|
|
db = self.connect()
|
|
data = convert(data)
|
|
if _id is not None:
|
|
data['_id'] = _id
|
|
_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
|
|
|
|
:returns: True if the document was already present
|
|
:rtype: bool
|
|
"""
|
|
db = self.connect()
|
|
data = 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, increment)
|
|
: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
|
|
|
|
:returns: a boolean (True if an existing document was updated) and the document after the update
|
|
:rtype: tuple of (bool, dict)
|
|
"""
|
|
db = self.connect()
|
|
data = convert(data or {})
|
|
_id_or_query = 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.items():
|
|
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 updateMany(self, collection, query, data):
|
|
"""Update multiple existing documents.
|
|
|
|
query can be an ID or a dict representing a query.
|
|
|
|
:param collection: update documents in this collection
|
|
:type collection: str
|
|
:param query: a query or a list of attributes in the data that must match
|
|
:type query: str or :class:`~bson.objectid.ObjectId` or iterable
|
|
:param data: the updated information to store
|
|
:type data: dict
|
|
|
|
:returns: a dict with the success state and number of updated items
|
|
:rtype: dict
|
|
"""
|
|
db = self.connect()
|
|
data = convert(data or {})
|
|
query = convert(query)
|
|
if not isinstance(query, dict):
|
|
query = {'_id': query}
|
|
if '_id' in data:
|
|
del data['_id']
|
|
return db[collection].update(query, {'$set': data}, multi=True)
|
|
|
|
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
|
|
|
|
:returns: dictionary with the number or removed documents
|
|
:rtype: dict
|
|
"""
|
|
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 = convert(_id_or_query)
|
|
return db[collection].remove(_id_or_query)
|