enable datastore emulator

jinfantes · Nov 5, 2019 · 2162ca0 · 2162ca0
1 parent 03e86e4
commit 2162ca0
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/flask_session/__init__.py b/flask_session/__init__.py
@@ -79,6 +79,8 @@ def _get_interface(self, app):
         config.setdefault('SESSION_SQLALCHEMY', None)
         config.setdefault('SESSION_SQLALCHEMY_TABLE', 'sessions')
 
+        config.setdefault('GCLOUD_APP_PROJECT_ID', 'unknown')
+
         if config['SESSION_TYPE'] == 'redis':
             session_interface = RedisSessionInterface(
                 config['SESSION_REDIS'], config['SESSION_KEY_PREFIX'],
@@ -106,8 +108,8 @@ def _get_interface(self, app):
                 config['SESSION_PERMANENT'])
         elif config['SESSION_TYPE'] == 'datastore':
             session_interface = GoogleCloudDatastoreSessionInterface(
-                config['SESSION_KEY_PREFIX'], config['SESSION_USE_SIGNER'],
-                config['SESSION_PERMANENT'])
+                config['GCLOUD_APP_PROJECT_ID'], config['SESSION_KEY_PREFIX'],
+                config['SESSION_USE_SIGNER'], config['SESSION_PERMANENT'])
         else:
             session_interface = NullSessionInterface()
 

diff --git a/flask_session/sessions.py b/flask_session/sessions.py
@@ -12,6 +12,7 @@
 import time
 import pytz
 import logging
+import os
 from datetime import datetime
 from uuid import uuid4
 try:
@@ -579,15 +580,26 @@ class GoogleCloudDatastoreSessionInterface(SessionInterface):
     serializer = pickle
     session_class = GoogleCloudDataStoreSession
 
-    def __init__(self, key_prefix, use_signer=False, permanent=True):
+    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.cloud import datastore
+        from google.auth import compute_engine
+        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):
         logging.warning('open_session')
         from google.cloud import datastore
-        ds_client = datastore.Client()
+        ds_client = self.get_client()
         sid = request.cookies.get(app.session_cookie_name)
         if not sid:
             sid = self._generate_sid()
@@ -620,7 +632,8 @@ def open_session(self, app, request):
 
     def save_session(self, app, session, response):
         logging.warning('save_session')
-        ds_client = datastore.Client()
+        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