#!/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('')

    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))