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