102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
from connection import _get_db
|
|
|
|
import pymongo
|
|
|
|
|
|
class QuerySet(object):
|
|
"""A set of results returned from a query. Wraps a MongoDB cursor,
|
|
providing Document objects as the results.
|
|
"""
|
|
|
|
def __init__(self, document, cursor):
|
|
self._document = document
|
|
self._cursor = cursor
|
|
|
|
def next(self):
|
|
"""Wrap the result in a Document object.
|
|
"""
|
|
return self._document._from_son(self._cursor.next())
|
|
|
|
def count(self):
|
|
"""Count the selected elements in the query.
|
|
"""
|
|
return self._cursor.count()
|
|
|
|
def limit(self, n):
|
|
"""Limit the number of returned documents to.
|
|
"""
|
|
self._cursor.limit(n)
|
|
# Return self to allow chaining
|
|
return self
|
|
|
|
def skip(self, n):
|
|
"""Skip n documents before returning the results.
|
|
"""
|
|
self._cursor.skip(n)
|
|
return self
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
|
|
class CollectionManager(object):
|
|
|
|
def __init__(self, document):
|
|
"""Set up the collection manager for a specific document.
|
|
"""
|
|
db = _get_db()
|
|
self._document = document
|
|
self._collection_name = document._meta['collection']
|
|
# This will create the collection if it doesn't exist
|
|
self._collection = db[self._collection_name]
|
|
|
|
def _save_document(self, document):
|
|
"""Save the provided document to the collection.
|
|
"""
|
|
_id = self._collection.save(document._to_mongo())
|
|
document._id = _id
|
|
|
|
def _transform_query(self, **query):
|
|
"""Transform a query from Django-style format to Mongo format.
|
|
"""
|
|
operators = ['neq', 'gt', 'gte', 'lt', 'lte', 'in', 'nin', 'mod',
|
|
'all', 'size', 'exists']
|
|
|
|
mongo_query = {}
|
|
for key, value in query.items():
|
|
parts = key.split('__')
|
|
# Check for an operator and transform to mongo-style if there is
|
|
if parts[-1] in operators:
|
|
op = parts.pop()
|
|
value = {'$' + op: value}
|
|
|
|
key = '.'.join(parts)
|
|
mongo_query[key] = value
|
|
|
|
return mongo_query
|
|
|
|
def find(self, **query):
|
|
"""Query the collection for documents matching the provided query.
|
|
"""
|
|
query = self._transform_query(**query)
|
|
query['_types'] = self._document._class_name
|
|
return QuerySet(self._document, self._collection.find(query))
|
|
|
|
def find_one(self, object_id=None, **query):
|
|
"""Query the collection for document matching the provided query.
|
|
"""
|
|
if object_id:
|
|
# Use just object_id if provided
|
|
if not isinstance(object_id, pymongo.objectid.ObjectId):
|
|
object_id = pymongo.objectid.ObjectId(object_id)
|
|
query = object_id
|
|
else:
|
|
# Otherwise, use the query provided
|
|
query = self._transform_query(**query)
|
|
query['_types'] = self._document._class_name
|
|
|
|
result = self._collection.find_one(query)
|
|
if result is not None:
|
|
result = self._document._from_son(result)
|
|
return result
|