1 files changed, 277 insertions, 0 deletions
diff --git a/google_appengine/google/appengine/tools/dev_appserver_index.py b/google_appengine/google/appengine/tools/dev_appserver_index.py
new file mode 100755
index 0000000..d69f656
--- /dev/null
+++ b/google_appengine/google/appengine/tools/dev_appserver_index.py
@@ -0,0 +1,277 @@
+#!/usr/bin/env python
+#
+# Copyright 2007 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Utilities for generating and updating index.yaml."""
+
+
+
+import os
+import logging
+
+from google.appengine.api import apiproxy_stub_map
+from google.appengine.api import datastore_admin
+from google.appengine.api import yaml_errors
+from google.appengine.datastore import datastore_index
+
+import yaml
+
+AUTO_MARKER = '\n# AUTOGENERATED\n'
+
+AUTO_COMMENT = '''
+# This index.yaml is automatically updated whenever the dev_appserver
+# detects that a new type of query is run.  If you want to manage the
+# index.yaml file manually, remove the above marker line (the line
+# saying "# AUTOGENERATED").  If you want to manage some indexes
+# manually, move them above the marker line.  The index.yaml file is
+# automatically uploaded to the admin console when you next deploy
+# your application using appcfg.py.
+'''
+
+
+def GenerateIndexFromHistory(query_history,
+                             all_indexes=None, manual_indexes=None):
+  """Generate most of the text for index.yaml from the query history.
+
+  Args:
+    query_history: Query history, a dict mapping query
+    all_indexes: Optional datastore_index.IndexDefinitions instance
+      representing all the indexes found in the input file.  May be None.
+    manual_indexes: Optional datastore_index.IndexDefinitions instance
+      containing indexes for which we should not generate output.  May be None.
+
+  Returns:
+    A string representation that can safely be appended to an
+    existing index.yaml file.
+  """
+
+  all_keys = datastore_index.IndexDefinitionsToKeys(all_indexes)
+  manual_keys = datastore_index.IndexDefinitionsToKeys(manual_indexes)
+
+  indexes = dict((key, 0) for key in all_keys - manual_keys)
+
+  for query, count in query_history.iteritems():
+    required, kind, ancestor, props, num_eq_filters = datastore_index.CompositeIndexForQuery(query)
+    if required:
+      key = (kind, ancestor, props)
+      if key not in manual_keys:
+        if key in indexes:
+          indexes[key] += count
+        else:
+          indexes[key] = count
+
+  res = []
+  for (kind, ancestor, props), count in sorted(indexes.iteritems()):
+    res.append('')
+    if count == 0:
+      message = '# Unused in query history -- copied from input.'
+    elif count == 1:
+      message = '# Used once in query history.'
+    else:
+      message = '# Used %d times in query history.' % count
+    res.append(message)
+    res.append(datastore_index.IndexYamlForQuery(kind, ancestor, props))
+
+  res.append('')
+  return '\n'.join(res)
+
+
+class IndexYamlUpdater(object):
+  """Helper class for updating index.yaml.
+
+  This class maintains some state about the query history and the
+  index.yaml file in order to minimize the number of times index.yaml
+  is actually overwritten.
+  """
+
+  index_yaml_is_manual = False
+  index_yaml_mtime = 0
+  last_history_size = 0
+
+  def __init__(self, root_path):
+    """Constructor.
+
+    Args:
+      root_path: Path to the app's root directory.
+    """
+    self.root_path = root_path
+
+  def UpdateIndexYaml(self, openfile=open):
+    """Update index.yaml.
+
+    Args:
+      openfile: Used for dependency injection.
+
+    We only ever write to index.yaml if either:
+    - it doesn't exist yet; or
+    - it contains an 'AUTOGENERATED' comment.
+
+    All indexes *before* the AUTOGENERATED comment will be written
+    back unchanged.  All indexes *after* the AUTOGENERATED comment
+    will be updated with the latest query counts (query counts are
+    reset by --clear_datastore).  Indexes that aren't yet in the file
+    will be appended to the AUTOGENERATED section.
+
+    We keep track of some data in order to avoid doing repetitive work:
+    - if index.yaml is fully manual, we keep track of its mtime to
+      avoid parsing it over and over;
+    - we keep track of the number of keys in the history dict since
+      the last time we updated index.yaml (or decided there was
+      nothing to update).
+    """
+    index_yaml_file = os.path.join(self.root_path, 'index.yaml')
+
+    try:
+      index_yaml_mtime = os.path.getmtime(index_yaml_file)
+    except os.error:
+      index_yaml_mtime = None
+
+    index_yaml_changed = (index_yaml_mtime != self.index_yaml_mtime)
+    self.index_yaml_mtime = index_yaml_mtime
+
+    datastore_stub = apiproxy_stub_map.apiproxy.GetStub('datastore_v3')
+    query_history = datastore_stub.QueryHistory()
+    history_changed = (len(query_history) != self.last_history_size)
+    self.last_history_size = len(query_history)
+
+    if not (index_yaml_changed or history_changed):
+      logging.debug('No need to update index.yaml')
+      return
+
+    if self.index_yaml_is_manual and not index_yaml_changed:
+        logging.debug('Will not update manual index.yaml')
+        return
+
+    if index_yaml_mtime is None:
+      index_yaml_data = None
+    else:
+      try:
+        fh = open(index_yaml_file, 'r')
+      except IOError:
+        index_yaml_data = None
+      else:
+        try:
+          index_yaml_data = fh.read()
+        finally:
+          fh.close()
+
+    self.index_yaml_is_manual = (index_yaml_data is not None and
+                                 AUTO_MARKER not in index_yaml_data)
+    if self.index_yaml_is_manual:
+      logging.info('Detected manual index.yaml, will not update')
+      return
+
+    if index_yaml_data is None:
+      all_indexes = None
+    else:
+      try:
+        all_indexes = datastore_index.ParseIndexDefinitions(index_yaml_data)
+      except yaml_errors.EventListenerError, e:
+        logging.error('Error parsing %s:\n%s', index_yaml_file, e)
+        return
+      except Exception, err:
+        logging.error('Error parsing %s:\n%s.%s: %s', index_yaml_file,
+                      err.__class__.__module__, err.__class__.__name__, err)
+        return
+
+    if index_yaml_data is None:
+      manual_part, automatic_part = 'indexes:\n', ''
+      manual_indexes = None
+    else:
+      manual_part, automatic_part = index_yaml_data.split(AUTO_MARKER, 1)
+      try:
+        manual_indexes = datastore_index.ParseIndexDefinitions(manual_part)
+      except Exception, err:
+        logging.error('Error parsing manual part of %s: %s',
+                      index_yaml_file, err)
+        return
+
+    automatic_part = GenerateIndexFromHistory(query_history,
+                                              all_indexes, manual_indexes)
+
+    try:
+      fh = openfile(index_yaml_file, 'w')
+    except IOError, err:
+      logging.error('Can\'t write index.yaml: %s', err)
+      return
+
+    try:
+      logging.info('Updating %s', index_yaml_file)
+      fh.write(manual_part)
+      fh.write(AUTO_MARKER)
+      fh.write(AUTO_COMMENT)
+      fh.write(automatic_part)
+    finally:
+      fh.close()
+
+    try:
+      self.index_yaml_mtime = os.path.getmtime(index_yaml_file)
+    except os.error, err:
+      logging.error('Can\'t stat index.yaml we just wrote: %s', err)
+      self.index_yaml_mtime = None
+
+
+def SetupIndexes(app_id, root_path):
+  """Ensure that the set of existing composite indexes matches index.yaml.
+
+  Note: this is similar to the algorithm used by the admin console for
+  the same purpose.
+
+  Args:
+    app_id: Application ID being served.
+    root_path: Path to the root of the application.
+  """
+  index_yaml_file = os.path.join(root_path, 'index.yaml')
+  try:
+    fh = open(index_yaml_file, 'r')
+  except IOError:
+    index_yaml_data = None
+  else:
+    try:
+      index_yaml_data = fh.read()
+    finally:
+      fh.close()
+
+  indexes = []
+  if index_yaml_data is not None:
+    index_defs = datastore_index.ParseIndexDefinitions(index_yaml_data)
+    if index_defs is not None:
+      indexes = index_defs.indexes
+      if indexes is None:
+        indexes = []
+
+  requested_indexes = datastore_admin.IndexDefinitionsToProtos(app_id, indexes)
+
+  existing_indexes = datastore_admin.GetIndices(app_id)
+
+  requested = dict((x.definition().Encode(), x) for x in requested_indexes)
+  existing = dict((x.definition().Encode(), x) for x in existing_indexes)
+
+  created = 0
+  for key, index in requested.iteritems():
+    if key not in existing:
+      datastore_admin.CreateIndex(index)
+      created += 1
+
+  deleted = 0
+  for key, index in existing.iteritems():
+    if key not in requested:
+      datastore_admin.DeleteIndex(index)
+      deleted += 1
+
+  if created or deleted:
+    logging.info("Created %d and deleted %d index(es); total %d",
+                 created, deleted, len(requested))