Docserver overhaul: Gitiles away from me.

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

 # Copyright 2014 The Chromium Authors. All rights reserved.
 # 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
+# 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'
 
-  /_enqueue/_refresh/content_providers/cr-native-client?commit=123ff65468dcafff0
-
-  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):
   '''Provides read access to the commit ID cache within the server. For example:
 
   /_query_commit/master
 
   will return the commit ID stored under the commit key "master" within the
   commit cache. Currently "master" is the only named commit we cache, and it
   corresponds to the commit ID whose data currently populates the data cache
@@ skipping 14 matching lines @@
           commit_id, history_log)
       return response
 
     commit_name = self._request.path
     id_future = commit_tracker.Get(commit_name)
     history_future = commit_tracker.GetHistory(commit_name)
     return Response.Ok(
         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!')