1 #!/usr/bin/env python

     2 #

     3 # Copyright 2007 Google Inc.

     4 #

     5 # Licensed under the Apache License, Version 2.0 (the "License");

     6 # you may not use this file except in compliance with the License.

     7 # You may obtain a copy of the License at

     8 #

     9 #     http://www.apache.org/licenses/LICENSE-2.0

    10 #

    11 # Unless required by applicable law or agreed to in writing, software

    12 # distributed under the License is distributed on an "AS IS" BASIS,

    13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

    14 # See the License for the specific language governing permissions and

    15 # limitations under the License.

    16 #

    17 

    18 """Utilities for generating and updating index.yaml."""

    19 

    20 

    21 

    22 import os

    23 import logging

    24 

    25 from google.appengine.api import apiproxy_stub_map

    26 from google.appengine.api import datastore_admin

    27 from google.appengine.api import yaml_errors

    28 from google.appengine.datastore import datastore_index

    29 

    30 import yaml

    31 

    32 AUTO_MARKER = '\n# AUTOGENERATED\n'

    33 

    34 AUTO_COMMENT = '''

    35 # This index.yaml is automatically updated whenever the dev_appserver

    36 # detects that a new type of query is run.  If you want to manage the

    37 # index.yaml file manually, remove the above marker line (the line

    38 # saying "# AUTOGENERATED").  If you want to manage some indexes

    39 # manually, move them above the marker line.  The index.yaml file is

    40 # automatically uploaded to the admin console when you next deploy

    41 # your application using appcfg.py.

    42 '''

    43 

    44 

    45 def GenerateIndexFromHistory(query_history,

    46                              all_indexes=None, manual_indexes=None):

    47   """Generate most of the text for index.yaml from the query history.

    48 

    49   Args:

    50     query_history: Query history, a dict mapping query

    51     all_indexes: Optional datastore_index.IndexDefinitions instance

    52       representing all the indexes found in the input file.  May be None.

    53     manual_indexes: Optional datastore_index.IndexDefinitions instance

    54       containing indexes for which we should not generate output.  May be None.

    55 

    56   Returns:

    57     A string representation that can safely be appended to an

    58     existing index.yaml file.

    59   """

    60 

    61   all_keys = datastore_index.IndexDefinitionsToKeys(all_indexes)

    62   manual_keys = datastore_index.IndexDefinitionsToKeys(manual_indexes)

    63 

    64   indexes = dict((key, 0) for key in all_keys - manual_keys)

    65 

    66   for query, count in query_history.iteritems():

    67     key = datastore_index.CompositeIndexForQuery(query)

    68     if key is not None:

    69       key = key[:3]

    70       if key not in manual_keys:

    71         if key in indexes:

    72           indexes[key] += count

    73         else:

    74           indexes[key] = count

    75 

    76   res = []

    77   for (kind, ancestor, props), count in sorted(indexes.iteritems()):

    78     res.append('')

    79     if count == 0:

    80       message = '# Unused in query history -- copied from input.'

    81     elif count == 1:

    82       message = '# Used once in query history.'

    83     else:

    84       message = '# Used %d times in query history.' % count

    85     res.append(message)

    86     res.append(datastore_index.IndexYamlForQuery(kind, ancestor, props))

    87 

    88   res.append('')

    89   return '\n'.join(res)

    90 

    91 

    92 class IndexYamlUpdater(object):

    93   """Helper class for updating index.yaml.

    94 

    95   This class maintains some state about the query history and the

    96   index.yaml file in order to minimize the number of times index.yaml

    97   is actually overwritten.

    98   """

    99 

   100   index_yaml_is_manual = False

   101   index_yaml_mtime = 0

   102   last_history_size = 0

   103 

   104   def __init__(self, root_path):

   105     """Constructor.

   106 

   107     Args:

   108       root_path: Path to the app's root directory.

   109     """

   110     self.root_path = root_path

   111 

   112   def UpdateIndexYaml(self, openfile=open):

   113     """Update index.yaml.

   114 

   115     Args:

   116       openfile: Used for dependency injection.

   117 

   118     We only ever write to index.yaml if either:

   119     - it doesn't exist yet; or

   120     - it contains an 'AUTOGENERATED' comment.

   121 

   122     All indexes *before* the AUTOGENERATED comment will be written

   123     back unchanged.  All indexes *after* the AUTOGENERATED comment

   124     will be updated with the latest query counts (query counts are

   125     reset by --clear_datastore).  Indexes that aren't yet in the file

   126     will be appended to the AUTOGENERATED section.

   127 

   128     We keep track of some data in order to avoid doing repetitive work:

   129     - if index.yaml is fully manual, we keep track of its mtime to

   130       avoid parsing it over and over;

   131     - we keep track of the number of keys in the history dict since

   132       the last time we updated index.yaml (or decided there was

   133       nothing to update).

   134     """

   135     index_yaml_file = os.path.join(self.root_path, 'index.yaml')

   136 

   137     try:

   138       index_yaml_mtime = os.path.getmtime(index_yaml_file)

   139     except os.error:

   140       index_yaml_mtime = None

   141 

   142     index_yaml_changed = (index_yaml_mtime != self.index_yaml_mtime)

   143     self.index_yaml_mtime = index_yaml_mtime

   144 

   145     datastore_stub = apiproxy_stub_map.apiproxy.GetStub('datastore_v3')

   146     query_history = datastore_stub.QueryHistory()

   147     history_changed = (len(query_history) != self.last_history_size)

   148     self.last_history_size = len(query_history)

   149 

   150     if not (index_yaml_changed or history_changed):

   151       logging.debug('No need to update index.yaml')

   152       return

   153 

   154     if self.index_yaml_is_manual and not index_yaml_changed:

   155         logging.debug('Will not update manual index.yaml')

   156         return

   157 

   158     if index_yaml_mtime is None:

   159       index_yaml_data = None

   160     else:

   161       try:

   162         fh = open(index_yaml_file, 'r')

   163       except IOError:

   164         index_yaml_data = None

   165       else:

   166         try:

   167           index_yaml_data = fh.read()

   168         finally:

   169           fh.close()

   170 

   171     self.index_yaml_is_manual = (index_yaml_data is not None and

   172                                  AUTO_MARKER not in index_yaml_data)

   173     if self.index_yaml_is_manual:

   174       logging.info('Detected manual index.yaml, will not update')

   175       return

   176 

   177     if index_yaml_data is None:

   178       all_indexes = None

   179     else:

   180       try:

   181         all_indexes = datastore_index.ParseIndexDefinitions(index_yaml_data)

   182       except yaml_errors.EventListenerError, e:

   183         logging.error('Error parsing %s:\n%s', index_yaml_file, e)

   184         return

   185       except Exception, err:

   186         logging.error('Error parsing %s:\n%s.%s: %s', index_yaml_file,

   187                       err.__class__.__module__, err.__class__.__name__, err)

   188         return

   189 

   190     if index_yaml_data is None:

   191       manual_part, automatic_part = 'indexes:\n', ''

   192       manual_indexes = None

   193     else:

   194       manual_part, automatic_part = index_yaml_data.split(AUTO_MARKER, 1)

   195       try:

   196         manual_indexes = datastore_index.ParseIndexDefinitions(manual_part)

   197       except Exception, err:

   198         logging.error('Error parsing manual part of %s: %s',

   199                       index_yaml_file, err)

   200         return

   201 

   202     automatic_part = GenerateIndexFromHistory(query_history,

   203                                               all_indexes, manual_indexes)

   204 

   205     try:

   206       fh = openfile(index_yaml_file, 'w')

   207     except IOError, err:

   208       logging.error('Can\'t write index.yaml: %s', err)

   209       return

   210 

   211     try:

   212       logging.info('Updating %s', index_yaml_file)

   213       fh.write(manual_part)

   214       fh.write(AUTO_MARKER)

   215       fh.write(AUTO_COMMENT)

   216       fh.write(automatic_part)

   217     finally:

   218       fh.close()

   219 

   220     try:

   221       self.index_yaml_mtime = os.path.getmtime(index_yaml_file)

   222     except os.error, err:

   223       logging.error('Can\'t stat index.yaml we just wrote: %s', err)

   224       self.index_yaml_mtime = None

   225 

   226 

   227 def SetupIndexes(app_id, root_path):

   228   """Ensure that the set of existing composite indexes matches index.yaml.

   229 

   230   Note: this is similar to the algorithm used by the admin console for

   231   the same purpose.

   232 

   233   Args:

   234     app_id: Application ID being served.

   235     root_path: Path to the root of the application.

   236   """

   237   index_yaml_file = os.path.join(root_path, 'index.yaml')

   238   try:

   239     fh = open(index_yaml_file, 'r')

   240   except IOError:

   241     index_yaml_data = None

   242   else:

   243     try:

   244       index_yaml_data = fh.read()

   245     finally:

   246       fh.close()

   247 

   248   indexes = []

   249   if index_yaml_data is not None:

   250     index_defs = datastore_index.ParseIndexDefinitions(index_yaml_data)

   251     if index_defs is not None:

   252       indexes = index_defs.indexes

   253       if indexes is None:

   254         indexes = []

   255 

   256   requested_indexes = datastore_admin.IndexDefinitionsToProtos(app_id, indexes)

   257 

   258   existing_indexes = datastore_admin.GetIndices(app_id)

   259 

   260   requested = dict((x.definition().Encode(), x) for x in requested_indexes)

   261   existing = dict((x.definition().Encode(), x) for x in existing_indexes)

   262 

   263   created = 0

   264   for key, index in requested.iteritems():

   265     if key not in existing:

   266       datastore_admin.CreateIndex(index)

   267       created += 1

   268 

   269   deleted = 0

   270   for key, index in existing.iteritems():

   271     if key not in requested:

   272       datastore_admin.DeleteIndex(index)

   273       deleted += 1

   274 

   275   if created or deleted:

   276     logging.info("Created %d and deleted %d index(es); total %d",

   277                  created, deleted, len(requested))