Docserver overhaul: Gitiles away from me.

chrome/common/extensions/docs/server2/admin_servlets.py

Index: chrome/common/extensions/docs/server2/admin_servlets.py
diff --git a/chrome/common/extensions/docs/server2/admin_servlets.py b/chrome/common/extensions/docs/server2/admin_servlets.py
index 7990761e02d1cb9928580904ccd657a0d13b4a44..4e0a838a8bcb86ce177e0508688297bed8aa7436 100644
--- a/chrome/common/extensions/docs/server2/admin_servlets.py
+++ b/chrome/common/extensions/docs/server2/admin_servlets.py
@@ -2,35 +2,28 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-from appengine_wrappers import taskqueue
+import json
+import logging
+
+from appengine_wrappers import memcache
 from commit_tracker import CommitTracker
+from environment import IsDevServer
+from environment_wrappers import CreateUrlFetcher
 from future import All
 from object_store_creator import ObjectStoreCreator
-from refresh_tracker import RefreshTracker
 from servlet import Servlet, Response
 
 
-class EnqueueServlet(Servlet):
-  '''This Servlet can be used to manually enqueue tasks on the default
-  taskqueue. Useful for when an admin wants to manually force a specific
-  DataSource refresh, but the refresh operation takes longer than the 60 sec
-  timeout of a non-taskqueue request. For example, you might query
-
-  /_enqueue/_refresh/content_providers/cr-native-client?commit=123ff65468dcafff0
+# This is the service account ID associated with the chrome-apps-doc project in
+# the Google Developers Console.
+_SERVICE_ACCOUNT_NAME = '636061184119-compute@developer.gserviceaccount.com'
 
-  which will enqueue a task (/_refresh/content_providers/cr-native-client) to
-  refresh the NaCl documentation cache for commit 123ff65468dcafff0.
-
-  Access to this servlet should always be restricted to administrative users.
-  '''
-  def __init__(self, request):
-    Servlet.__init__(self, request)
-
-  def Get(self):
-    queue = taskqueue.Queue()
-    queue.add(taskqueue.Task(url='/%s' % self._request.path,
-                             params=self._request.arguments))
-    return Response.Ok('Task enqueued.')
+# This is Google's OAuth2 service for retrieving user information from an access
+# token. It is used to authenticate an incoming access token in the metadata
+# flush servlet, ensuring that only requests from the above service account are
+# fulfilled.
+_ACCOUNT_INFO_URL = ('https://www.googleapis.com/oauth2/v1/userinfo?'
+                     'access_token=%s')
 
 
 class QueryCommitServlet(Servlet):
@@ -65,47 +58,78 @@ class QueryCommitServlet(Servlet):
         All((id_future, history_future)).Then(generate_response).Get())
 
 
-class DumpRefreshServlet(Servlet):
-  def __init__(self, request):
-    Servlet.__init__(self, request)
+class FlushMemcacheServlet(Servlet):
+  '''Flushes the entire memcache.
 
-  def Get(self):
-    object_store_creator = ObjectStoreCreator(start_empty=False)
-    refresh_tracker = RefreshTracker(object_store_creator)
-    commit_id = self._request.path
-    work_order = refresh_tracker._GetWorkOrder(commit_id).Get()
-    task_names = ['%s@%s' % (commit_id, task) for task in work_order.tasks]
-    completions = refresh_tracker._task_completions.GetMulti(task_names).Get()
-    missing = []
-    for task in task_names:
-      if task not in completions:
-        missing.append(task)
-    response = 'Missing:<br>%s' % ''.join('%s<br>' % task for task in missing)
-    return Response.Ok(response)
-
-class ResetCommitServlet(Servlet):
-  '''Writes a new commit ID to the commit cache. For example:
-
-  /_reset_commit/master/123456
-
-  will reset the 'master' commit ID to '123456'. The provided commit MUST be
-  in the named commit's recent history or it will be ignored.
+  This requires an access token for the project's main service account. Without
+  said token, the request is considered invalid.
   '''
 
   class Delegate(object):
-    def CreateCommitTracker(self):
-      return CommitTracker(ObjectStoreCreator(start_empty=False))
+    def IsAuthorized(self, access_token):
+      '''Verifies that a given access_token represents the main service account.
+      '''
+      fetcher = CreateUrlFetcher()
+      response = fetcher.Fetch(_ACCOUNT_INFO_URL % access_token)
+      if response.status_code != 200:
+        return False
+      try:
+        info = json.loads(response.content)
+      except:
+        return False
+      return info['email'] == _SERVICE_ACCOUNT_NAME
 
   def __init__(self, request, delegate=Delegate()):
     Servlet.__init__(self, request)
     self._delegate = delegate
 
+  def GetAccessToken(self):
+    auth_header = self._request.headers.get('Authorization')
+    if not auth_header:
+      return None
+    try:
+      method, token = auth_header.split(' ', 1)
+    except:
+      return None
+    if method != 'Bearer':
+      return None
+    return token
+
   def Get(self):
-    commit_tracker = self._delegate.CreateCommitTracker()
-    commit_name, commit_id = self._request.path.split('/', 1)
-    history = commit_tracker.GetHistory(commit_name).Get()
-    if not any(entry.commit_id == commit_id for entry in history):
-      return Response.BadRequest('Commit %s not cached.' % commit_id)
-    commit_tracker.Set(commit_name, commit_id).Get()
-    return Response.Ok('Commit "%s" updated to %s' % (commit_name, commit_id))
+    access_token = self.GetAccessToken()
+    if not access_token:
+      return Response.Unauthorized('Unauthorized', 'Bearer', 'update')
+    if not self._delegate.IsAuthorized(access_token):
+      return Response.Forbidden('Forbidden')
+    result = memcache.flush_all()
+    return Response.Ok('Flushed: %s' % result)
+
+
+class UpdateCacheServlet(Servlet):
+  '''Devserver-only servlet for pushing local file data into the datastore.
+  This is useful if you've used update_cache.py to build a local datastore
+  for testing. Query:
+
+      /_update_cache/FOO_DATA
 
+  to make the devserver read FOO_DATA from its pwd and push all the data into
+  datastore.
+  '''
+  def __init__(self, request):
+    Servlet.__init__(self, request)
+
+  def Get(self):
+    if not IsDevServer():
+      return Response.BadRequest('')
+    import cPickle
+    from persistent_object_store_appengine import PersistentObjectStoreAppengine
+    with open(self._request.path, 'r') as f:
+      data = cPickle.load(f)
+    for namespace, contents in data.iteritems():
+      store = PersistentObjectStoreAppengine(namespace)
+      for k, v in cPickle.loads(contents).iteritems():
+        try:
+          store.Set(k, v).Get()
+        except:
+          logging.warn('Skipping entry %s because of errors.' % k)
+    return Response.Ok('Data pushed!')