1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180
|
# Cork - Authentication module for the Bottle web framework
# Copyright (C) 2013 Federico Ceratto and others, see AUTHORS file.
# Released under LGPLv3+ license, see LICENSE.txt
"""
.. module:: mongodb_backend
:synopsis: MongoDB storage backend.
"""
from logging import getLogger
log = getLogger(__name__)
from .base_backend import Backend, Table
try:
import pymongo
is_pymongo_2 = (pymongo.version_tuple[0] == 2)
except ImportError: # pragma: no cover
pass
class MongoTable(Table):
"""Abstract MongoDB Table.
Allow dictionary-like access.
"""
def __init__(self, name, key_name, collection):
self._name = name
self._key_name = key_name
self._coll = collection
def create_index(self):
"""Create collection index."""
self._coll.create_index(
self._key_name,
drop_dups=True,
unique=True,
)
def __len__(self):
return self._coll.count()
def __contains__(self, value):
r = self._coll.find_one({self._key_name: value})
return r is not None
def __iter__(self):
"""Iter on dictionary keys"""
if is_pymongo_2:
r = self._coll.find(fields=[self._key_name,])
else:
r = self._coll.find(projection=[self._key_name,])
return (i[self._key_name] for i in r)
def iteritems(self):
"""Iter on dictionary items.
:returns: generator of (key, value) tuples
"""
r = self._coll.find()
for i in r:
d = i.copy()
d.pop(self._key_name)
d.pop('_id')
yield (i[self._key_name], d)
def pop(self, key_val):
"""Remove a dictionary item"""
r = self[key_val]
self._coll.remove({self._key_name: key_val}, w=1)
return r
class MongoSingleValueTable(MongoTable):
"""MongoDB table accessible as a simple key -> value dictionary.
Used to store roles.
"""
# Values are stored in a MongoDB "column" named "val"
def __init__(self, *args, **kw):
super(MongoSingleValueTable, self).__init__(*args, **kw)
def __setitem__(self, key_val, data):
assert not isinstance(data, dict)
spec = {self._key_name: key_val}
data = {self._key_name: key_val, 'val': data}
if is_pymongo_2:
self._coll.update(spec, {'$set': data}, upsert=True, w=1)
else:
self._coll.update_one(spec, {'$set': data}, upsert=True)
def __getitem__(self, key_val):
r = self._coll.find_one({self._key_name: key_val})
if r is None:
raise KeyError(key_val)
return r['val']
class MongoMutableDict(dict):
"""Represent an item from a Table. Acts as a dictionary.
"""
def __init__(self, parent, root_key, d):
"""Create a MongoMutableDict instance.
:param parent: Table instance
:type parent: :class:`MongoTable`
"""
super(MongoMutableDict, self).__init__(d)
self._parent = parent
self._root_key = root_key
def __setitem__(self, k, v):
super(MongoMutableDict, self).__setitem__(k, v)
spec = {self._parent._key_name: self._root_key}
if is_pymongo_2:
r = self._parent._coll.update(spec, {'$set': {k: v}}, upsert=True)
else:
r = self._parent._coll.update_one(spec, {'$set': {k: v}}, upsert=True)
class MongoMultiValueTable(MongoTable):
"""MongoDB table accessible as a dictionary.
"""
def __init__(self, *args, **kw):
super(MongoMultiValueTable, self).__init__(*args, **kw)
def __setitem__(self, key_val, data):
assert isinstance(data, dict)
key_name = self._key_name
if key_name in data:
assert data[key_name] == key_val
else:
data[key_name] = key_val
spec = {key_name: key_val}
if u'_id' in data:
del(data[u'_id'])
if is_pymongo_2:
self._coll.update(spec, {'$set': data}, upsert=True, w=1)
else:
self._coll.update_one(spec, {'$set': data}, upsert=True)
def __getitem__(self, key_val):
r = self._coll.find_one({self._key_name: key_val})
if r is None:
raise KeyError(key_val)
return MongoMutableDict(self, key_val, r)
class MongoDBBackend(Backend):
def __init__(self, db_name='cork', hostname='localhost', port=27017, initialize=False, username=None, password=None):
"""Initialize MongoDB Backend"""
connection = pymongo.MongoClient(host=hostname, port=port)
db = connection[db_name]
if username and password:
db.authenticate(username, password)
self.users = MongoMultiValueTable('users', 'login', db.users)
self.pending_registrations = MongoMultiValueTable(
'pending_registrations',
'pending_registration',
db.pending_registrations
)
self.roles = MongoSingleValueTable('roles', 'role', db.roles)
if initialize:
self._initialize_storage()
def _initialize_storage(self):
"""Create MongoDB indexes."""
for c in (self.users, self.roles, self.pending_registrations):
c.create_index()
def save_users(self):
pass
def save_roles(self):
pass
def save_pending_registrations(self):
pass
|