1# -*- coding: utf-8 -*-
2"""
3    webapp2_extras.appengine.sessions_ndb
4    =====================================
5
6    Extended sessions stored in datastore using the ndb library.
7
8    :copyright: 2011 by tipfy.org.
9    :license: Apache Sotware License, see LICENSE for details.
10"""
11from __future__ import absolute_import
12
13from google.appengine.api import memcache
14
15try:
16    from ndb import model
17except ImportError: # pragma: no cover
18    from google.appengine.ext.ndb import model
19
20try:
21    from ndb.model import PickleProperty
22except ImportError: # pragma: no cover
23    try:
24        from google.appengine.ext.ndb.model import PickleProperty
25    except ImportError: # pragma: no cover
26        # ndb in SDK 1.6.1 doesn't have PickleProperty.
27        import pickle
28
29        class PickleProperty(model.BlobProperty):
30            """A Property whose value is any picklable Python object."""
31
32            def _validate(self, value):
33                return value
34
35            def _db_set_value(self, v, p, value):
36                super(PickleProperty, self)._db_set_value(v, p,
37                    pickle.dumps(value))
38
39            def _db_get_value(self, v, p):
40                if not v.has_stringvalue():
41                    return None
42
43                return pickle.loads(v.stringvalue())
44
45
46from webapp2_extras import sessions
47
48class Session(model.Model):
49    """A model to store session data."""
50
51    #: Save time.
52    updated = model.DateTimeProperty(auto_now=True)
53    #: Session data, pickled.
54    data = PickleProperty()
55
56    @classmethod
57    def get_by_sid(cls, sid):
58        """Returns a ``Session`` instance by session id.
59
60        :param sid:
61            A session id.
62        :returns:
63            An existing ``Session`` entity.
64        """
65        data = memcache.get(sid)
66        if not data:
67            session = model.Key(cls, sid).get()
68            if session:
69                data = session.data
70                memcache.set(sid, data)
71
72        return data
73
74    def _put(self):
75        """Saves the session and updates the memcache entry."""
76        memcache.set(self._key.id(), self.data)
77        super(Session, self).put()
78
79
80class DatastoreSessionFactory(sessions.CustomBackendSessionFactory):
81    """A session factory that stores data serialized in datastore.
82
83    To use datastore sessions, pass this class as the `factory` keyword to
84    :meth:`webapp2_extras.sessions.SessionStore.get_session`::
85
86        from webapp2_extras import sessions_ndb
87
88        # [...]
89
90        session = self.session_store.get_session(
91            name='db_session', factory=sessions_ndb.DatastoreSessionFactory)
92
93    See in :meth:`webapp2_extras.sessions.SessionStore` an example of how to
94    make sessions available in a :class:`webapp2.RequestHandler`.
95    """
96
97    #: The session model class.
98    session_model = Session
99
100    def _get_by_sid(self, sid):
101        """Returns a session given a session id."""
102        if self._is_valid_sid(sid):
103            data = self.session_model.get_by_sid(sid)
104            if data is not None:
105                self.sid = sid
106                return sessions.SessionDict(self, data=data)
107
108        self.sid = self._get_new_sid()
109        return sessions.SessionDict(self, new=True)
110
111    def save_session(self, response):
112        if self.session is None or not self.session.modified:
113            return
114
115        self.session_model(id=self.sid, data=dict(self.session))._put()
116        self.session_store.save_secure_cookie(
117            response, self.name, {'_sid': self.sid}, **self.session_args)
118