Add ndb support.

gaesessions/__init__.py

@@ -9,9 +9,11 @@
 import os
 import threading
 import time
+import traceback
 
 from google.appengine.api import memcache
-from google.appengine.ext import db
+from google.appengine.ext import db, ndb
+from google.appengine.datastore import entity_pb
 
 # Configurable cookie options
 COOKIE_NAME_PREFIX = "DgU"  # identifies a cookie as being one used by gae-sessions (so you can set cookies too)
@@ -47,10 +49,14 @@ def is_gaesessions_key(k):
     return k.startswith(COOKIE_NAME_PREFIX)
 
 
-class SessionModel(db.Model):
+class SessionModel(ndb.Model):
     """Contains session data.  key_name is the session ID and pdump contains a
     pickled dictionary which maps session variables to their values."""
-    pdump = db.BlobProperty()
+    pdump = ndb.BlobProperty()
+
+    @classmethod
+    def _get_kind(cls):
+        return 'SeS'
 
 
 class Session(object):
@@ -59,7 +65,7 @@ class Session(object):
     ``sid`` - if set, then the session for that sid (if any) is loaded. Otherwise,
     sid will be loaded from the HTTP_COOKIE (if any).
     """
-    DIRTY_BUT_DONT_PERSIST_TO_DB = 1
+    DIRTY_BUT_DONT_PERSIST_xTO_DB = 1
 
     def __init__(self, sid=None, lifetime=DEFAULT_LIFETIME, no_datastore=False,
                  cookie_only_threshold=DEFAULT_COOKIE_ONLY_THRESH, cookie_key=None):
@@ -88,32 +94,50 @@ def __compute_hmac(base_key, sid, text):
         return b64encode(hmac.new(key, text, hashlib.sha256).digest())
 
     def __read_cookie(self):
+        if 'HTTP_COOKIE' in os.environ:
+            self.load_cookie(os.environ['HTTP_COOKIE'])
+
+    def load_cookie(self, cookie_string):
         """Reads the HTTP Cookie and loads the sid and data from it (if any)."""
         try:
             # check the cookie to see if a session has been started
-            cookie = SimpleCookie(os.environ['HTTP_COOKIE'])
+            cookie = SimpleCookie(cookie_string)
+            logging.info("JUST GOT COOKIE WHICH IS: {}".format(cookie))
             self.cookie_keys = filter(is_gaesessions_key, cookie.keys())
             if not self.cookie_keys:
+                logging.info ("NO APPROPRIATE KEYS")
                 return  # no session yet
             self.cookie_keys.sort()
+            logging.info ("THE KEYS: {}".format(self.cookie_keys))
             data = ''.join(cookie[k].value for k in self.cookie_keys)
+            logging.info("DATA: {}".format(data))
             i = SIG_LEN + SID_LEN
+            logging.info("i: {}".format(i))
             sig, sid, b64pdump = data[:SIG_LEN], data[SIG_LEN:i], data[i:]
+            logging.info("sig: {}".format(sig))
+            logging.info("sid: {}".format(sid))
+            logging.info("b64pdump: {}".format(b64pdump))
             pdump = b64decode(b64pdump)
             actual_sig = Session.__compute_hmac(self.base_key, sid, pdump)
             if sig == actual_sig:
+                logging.info("SIG MATCH")
                 self.__set_sid(sid, False)
                 # check for expiration and terminate the session if it has expired
                 if self.get_expiration() != 0 and time.time() > self.get_expiration():
+                    logging.info("EXPIRED")
                     return self.terminate()
 
                 if pdump:
+                    logging.info("LOADING NEW DATA")
                     self.data = self.__decode_data(pdump)
                 else:
+                    logging.info("SETTING NONE DATA")
                     self.data = None  # data is in memcache/db: load it on-demand
             else:
                 logging.warn('cookie with invalid sig received from %s: %s' % (os.environ.get('REMOTE_ADDR'), b64pdump))
-        except (CookieError, KeyError, IndexError, TypeError):
+        except (CookieError, KeyError, IndexError, TypeError), e:
+            logging.info ("OOPSIE: {}".format(e))
+            logging.info(traceback.print_exc(e))
             # there is no cookie (i.e., no session) or the cookie is invalid
             self.terminate(False)
 
@@ -198,7 +222,9 @@ def __encode_data(d):
         eP = {}  # for models encoded as protobufs
         eO = {}  # for everything else
         for k, v in d.iteritems():
-            if isinstance(v, db.Model):
+            if isinstance(v, ndb.Model):
+                eP[k] = ndb.ModelAdapter().entity_to_pb(v).Encode()
+            elif isinstance(v, db.Model):
                 eP[k] = db.model_to_protobuf(v)
             else:
                 eO[k] = v
@@ -210,9 +236,15 @@ def __decode_data(pdump):
         try:
             eP, eO = pickle.loads(pdump)
             for k, v in eP.iteritems():
-                eO[k] = db.model_from_protobuf(v)
+                try:
+                    eO[k] = ndb.ModelAdapter().pb_to_entity(entity_pb.EntityProto(v))
+                except:
+                    try:
+                        eO[k] = db.model_from_protobuf(v)
+                    except Exception, e:
+                        logging.warn("failed ({}) to decode session key {}: {}".format(e.message, k))
         except Exception, e:
-            logging.warn("failed to decode session data: %s" % e)
+            logging.warn("Big failure decoding session data: %s" % e)
             eO = {}
         return eO
 
@@ -268,7 +300,7 @@ def __set_sid(self, sid, make_cookie=True):
         if self.sid:
             self.__clear_data()
         self.sid = sid
-        self.db_key = db.Key.from_path(SessionModel.kind(), sid, namespace='')
+        self.db_key = ndb.Key(SessionModel._get_kind(), sid, namespace='')
 
         # set the cookie if requested
         if make_cookie:
@@ -279,7 +311,7 @@ def __clear_data(self):
         if self.sid:
             memcache.delete(self.sid, namespace='')  # not really needed; it'll go away on its own
             try:
-                db.delete(self.db_key)
+                self.db_key.delete()
             except:
                 pass  # either it wasn't in the db (maybe cookie/memcache-only) or db is down => cron will expire it
 
@@ -294,7 +326,7 @@ def __retrieve_data(self):
                 logging.info("can't find session data in memcache for sid=%s (using memcache only sessions)" % self.sid)
                 self.terminate(False)  # we lost it; just kill the session
                 return
-            session_model_instance = db.get(self.db_key)
+            session_model_instance = self.db_key.get()
             if session_model_instance:
                 pdump = session_model_instance.pdump
             else:
@@ -341,7 +373,7 @@ def save(self, persist_even_if_using_cookie=False):
         if dirty is Session.DIRTY_BUT_DONT_PERSIST_TO_DB or self.no_datastore:
             return
         try:
-            SessionModel(key_name=self.sid, pdump=pdump).put()
+            SessionModel(id=self.sid, pdump=pdump).put()
         except Exception, e:
             logging.warning("unable to persist session to datastore for sid=%s (%s)" % (self.sid, e))
 
@@ -506,10 +538,16 @@ def delete_expired_sessions():
     Returns True if all expired sessions have been removed.
     """
     now_str = unicode(int(time.time()))
-    q = db.Query(SessionModel, keys_only=True, namespace='')
-    key = db.Key.from_path('SessionModel', now_str + u'\ufffd', namespace='')
-    q.filter('__key__ < ', key)
-    results = q.fetch(500)
-    db.delete(results)
-    logging.info('gae-sessions: deleted %d expired sessions from the datastore' % len(results))
-    return len(results) < 500
+    key = ndb.Key(SessionModel._get_kind(), now_str + u'\ufffd', namespace='')
+    logging.info ("FENCEPOST KEY: {}".format(key))
+    q = SessionModel.query(SessionModel.key < key, default_options=ndb.QueryOptions(keys_only=True, limit=500), namespace='')
+    class Mapper (object):
+        total = 0
+        def __call__ (self, key):
+            self.total += 1
+            #key.delete ()
+            logging.info("Would delete: {}".format(key))
+    m = Mapper ()
+    q.map(m)
+    logging.info('gae-sessions: deleted %d expired session(s) from the datastore' % m.total)
+    return m.total < 500