feat(google cloud datastore): added google cloud datastore as a sessi…

…on backend

pallets-eco#109

README.md

@@ -162,6 +162,12 @@ Uses elasticsearch as a session backend. ([elasticsearch](https://elasticsearch-
 - SESSION_ELASTICSEARCH_HOST
 - SESSION_ELASTICSEARCH_INDEX
 
+### `GoogleCloudDatastoreSessionInterface`
+
+Uses Google Cloud Datastore as a session backend. ([google-cloud-datastore](https://github.com/googleapis/python-datastore) required)
+
+- GCLOUD_APP_PROJECT_ID
+
 ## Credits
 
 This project is a fork of [flask-session](https://github.com/fengsp/flask-session), created by [Shipeng Feng](https://github.com/fengsp).

flask_session/__init__.py

@@ -15,6 +15,7 @@
 from .sessions import (
     ElasticsearchSessionInterface,
     FileSystemSessionInterface,
+    GoogleCloudDatastoreSessionInterface,
     MemcachedSessionInterface,
     MongoDBSessionInterface,
     NullSessionInterface,
@@ -96,6 +97,7 @@ def _get_interface(self, app):
         config.setdefault("SESSION_SQLALCHEMY", None)
         config.setdefault("SESSION_SQLALCHEMY_TABLE", "sessions")
         config.setdefault("SESSION_SQLALCHEMY_SEQUENCE", None)
+        config.setdefault("GCLOUD_APP_PROJECT_ID", "unknown")
 
         if config["SESSION_TYPE"] == "redis":
             session_interface = RedisSessionInterface(
@@ -150,6 +152,13 @@ def _get_interface(self, app):
                 config["SESSION_USE_SIGNER"],
                 config["SESSION_PERMANENT"],
             )
+        elif config["SESSION_TYPE"] == "datastore":
+            session_interface = GoogleCloudDatastoreSessionInterface(
+                config["GCLOUD_APP_PROJECT_ID"],
+                config["SESSION_KEY_PREFIX"],
+                config["SESSION_USE_SIGNER"],
+                config["SESSION_PERMANENT"],
+            )
         else:
             session_interface = NullSessionInterface()
 

flask_session/sessions.py

@@ -8,6 +8,7 @@
     :license: BSD, see LICENSE for more details.
 """
 
+import os
 import sys
 import time
 from datetime import datetime
@@ -74,6 +75,10 @@ class SqlAlchemySession(ServerSideSession):
     pass
 
 
+class GoogleCloudDataStoreSession(ServerSideSession):
+    pass
+
+
 class SessionInterface(FlaskSessionInterface):
     def _generate_sid(self):
         return str(uuid4())
@@ -800,3 +805,116 @@ def save_session(self, app, session, response):
             path=path,
             secure=secure,
         )
+
+
+class GoogleCloudDatastoreSessionInterface(SessionInterface):
+    """Uses the Google cloud datastore as a session backend.
+
+    :param key_prefix: A prefix that is added to all store keys.
+    :param use_signer: Whether to sign the session id cookie or not.
+    :param permanent: Whether to use permanent session or not.
+    """
+
+    serializer = pickle
+    session_class = GoogleCloudDataStoreSession
+
+    def __init__(self, gcloud_project, key_prefix, use_signer=False, permanent=True):
+        self.gcloud_project = gcloud_project
+        self.key_prefix = key_prefix
+        self.use_signer = use_signer
+        self.permanent = permanent
+
+    def get_client(self):
+        import requests
+        from google.auth import compute_engine
+        from google.cloud import datastore
+
+        if os.environ.get("DATASTORE_EMULATOR_HOST"):
+            return datastore.Client(
+                _http=requests.Session, project="virustotal-avs-control"
+            )
+        return datastore.Client(credentials=compute_engine.Credentials())
+
+    def open_session(self, app, request):
+        ds_client = self.get_client()
+        sid = request.cookies.get(app.config["SESSION_COOKIE_NAME"])
+        if not sid:
+            sid = self._generate_sid()
+            return self.session_class(sid=sid, permanent=self.permanent)
+        if self.use_signer:
+            signer = self._get_signer(app)
+            if signer is None:
+                return None
+            try:
+                sid_as_bytes = signer.unsign(sid)
+                sid = sid_as_bytes.decode()
+            except BadSignature:
+                sid = self._generate_sid()
+                return self.session_class(sid=sid, permanent=self.permanent)
+
+        store_id = self.key_prefix + sid
+        session_key = ds_client.key("session", store_id)
+        saved_session = ds_client.get(session_key)
+        if saved_session and saved_session["expiry"] <= pytz.utc.localize(
+            datetime.now()
+        ):
+            ds_client.delete(session_key)
+            saved_session = None
+        if saved_session:
+            try:
+                value = saved_session["data"]
+                data = self.serializer.loads(want_bytes(value))
+                return self.session_class(data, sid=sid)
+            except:
+                return self.session_class(sid=sid, permanent=self.permanent)
+        return self.session_class(sid=sid, permanent=self.permanent)
+
+    def save_session(self, app, session, response):
+        from google.cloud import datastore
+
+        ds_client = self.get_client()
+        domain = self.get_cookie_domain(app)
+        path = self.get_cookie_path(app)
+        store_id = self.key_prefix + session.sid
+        session_key = ds_client.key("session", store_id)
+        saved_session = ds_client.get(session_key)
+        if not session:
+            if session.modified:
+                if saved_session:
+                    ds_client.delete(session_key)
+                response.delete_cookie(
+                    app.config["SESSION_COOKIE_NAME"], domain=domain, path=path
+                )
+            return
+
+        httponly = self.get_cookie_httponly(app)
+        secure = self.get_cookie_secure(app)
+        expires = self.get_expiration_time(app, session)
+        value = self.serializer.dumps(dict(session))
+        if saved_session:
+            if not expires:
+                ds_client.delete(session_key)
+                return
+            saved_session["data"] = value
+            saved_session["expiry"] = expires
+            ds_client.put(saved_session)
+        else:
+            new_session = datastore.Entity(
+                key=session_key, exclude_from_indexes=("data",)
+            )
+            new_session["data"] = value
+            new_session["expiry"] = expires
+            ds_client.put(new_session)
+        if self.use_signer:
+            session_id = self._get_signer(app).sign(want_bytes(session.sid))
+        else:
+            session_id = session.sid
+        response.set_cookie(
+            app.config["SESSION_COOKIE_NAME"],
+            session_id,
+            expires=expires,
+            httponly=httponly,
+            domain=domain,
+            path=path,
+            secure=secure,
+        )