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