#!/usr/bin/env python
#
# Copyright 2010 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.


"""This module contains datastore trigger system implementation.

Overview
========

The system is designed to execute user's code every time an entity of specific
kind is changed by application. Right now the only execution model is delayed
execution (via taskqueue), which makes it useful only in eventually-consistent
scenarios.

It is important to remember that:

* Triggers won't be fired if you change your data not from the app itself (e.g.
  admin console).
* There might be a delay (sometimes even minute) between actual data change and
  trigger execution.
* Trigger executions for a single entity are batched. I.e. if you change entity
  quickly multiple times you'll most likely get your trigger function called
  only once.

Current Limitations:

* Right now only cron-based updates processing is implemented. This means that
  the delay between event and trigger processing might be about 1 minute in
  worst case.
* Due to limitations of current datastore api batch deletes are executed
  by triggers module as separate transactions. This may decrease performance
  of deleting code quite a lot.
* It is estimated that with this design you won't be able to process more than
  60 changes per second. See the implementation details and planned improvements
  sections for more details.


Usage
=====

To use triggers:

* Add the following code to your app.yaml::

    includes:
    - afterburner/experimental/db/triggers.yaml

* Add following lines to your cron.yaml::

    cron:
    - description: triggers processing
      url: /_ab/triggers/process
      schedule: every 1 minutes

* Create appengine_config.py in application root directory and import your
  modules with all trigger definitions from it.
* Define triggers as following::

    from afterburner.experimental.db import triggers
    @triggers.trigger(kind=Order.kind())
    def order_trigger(key):
      logging.debug("Order was changed: %s", key)

  Trigger function only argument is changed entity key. Because trigger execution
  is decoupled in time from actual db operation, the information about operation
  is not provided.
* Make sure that you're using run_wsgi_app() function to run your app and not
  older approach. For details check appengine documentation:
  http://code.google.com/appengine/docs/python/tools/webapp/running.html


Implementation Details
======================

To implement transactionally-safe triggers for every entity change we detect,
we create a child "Update" entity, which is later processed by updater
cron job. Updater cron job deletes Update entities and queues tasks for executing
triggers for particular entities.

Adding child entities to puts doesn't change runtime behavior of the system
at all, because datastore internally groups all entities in batch puts into
entity groups and performs their write together in one transaction automatically.

Unfortunately it's currently not possible to use the same approach to batch
deletes. We disassemble batch delete and delete each entity individually in
its own transaction together with creating an update entity.
Ideally datastore should provide an api where you can batch all kinds of
operations together.


Alternative Approaches
======================

This section outlines alternative approaches to triggers system designs. Some
of these approaches can augment already existing ones to improve performance
and other characteristics.

Transactional Tasks
-------------------

Instead of creating update tasks it's possible to use transactional tasks
to queue triggers processing for each entity. This works perfectly for
writes in transaction, but doesn't work good in case of batch operations
across multiple entity groups: you have to disassemble both puts and gets
the same way we disassemble deletes right now.

The benefit of this approach is that by eliminating Update entity and
update processor, trigger latency can be made quite low.

Another caveat with this approach is task's 10K payload size limit and limit
of 5 transactional tasks per transaction. If there are lots of entities in a
single transaction, it might be difficult to store keys for all of them
in 5 tasks. However, it looks like this limitation is not a problem for most
apps.

Immediate Triggers Execution
----------------------------

It was a specific goal for this library to implement delayed triggers execution.
Immediate trigger execution can be significant utility for certain use cases.


Future Improvements
===================

Following improvements to the library are considered:

* Use transactional tasks for triggers enqueing in situtations where it's
  possible (non-batched puts/deletes, small transactions)
* Optional unbatched trigger execution
* Augment current cron execution with chained tasks. This will allow bringing
  trigger latency down.

One of these can be an excellent project for beginning contributors.
"""

__author__ = 'mike.aizatsky@gmail.com (Mike Aizatsky)'


import logging
import urllib

from afterburner import util
from google.appengine.api import apiproxy_stub_map
from google.appengine.api import datastore
from google.appengine.api import taskqueue
from google.appengine.datastore import datastore_rpc
from google.appengine.ext import db
from google.appengine.ext import webapp


PROCESS_URL = "/_ab/triggers/process"
UPDATE_URL = "/_ab/triggers/update"


def _allocate_ids(keypb, size=1):
  """Version of db.allocate_ids that works inside the transaction.

  Args:
    keypb: Key as a protocol buffer.
    size: number of keys to allocate.
  """
  conn = datastore_rpc.Connection(adapter=None)
  return conn.allocate_ids(keypb, size)


class Update(db.Model):
  @classmethod
  def kind(cls):
    return "_ab_triggers_Update"


class TriggerManager(object):
  # entity kind => [trigger functions]
  triggers = {}
  ignore_deletes = False
  initialized = False

  @classmethod
  def pre_hook(cls, service, call, request, response):
    if call == "Put":
      extra_entities = []
      for entity in request.entity_list():
        key = entity.mutable_key()
        path = key.mutable_path()
        last_path_element = path.element_list()[-1]
        if (last_path_element.type() == Update.kind()
            or last_path_element.type() not in cls.triggers):
          # no triggers for this entity kind.
          continue
        if not last_path_element.id() and not last_path_element.has_name():
          # assign id to the entity
          last_path_element.set_id(-1)
          allocated_id = _allocate_ids(key)[0]
          last_path_element.set_id(allocated_id)

          # Populate entity group, because datastore stub is stupid.
          # Actual datastore doesn't need it, but it won't hurt anyway.
          group = entity.mutable_entity_group()
          root = entity.key().path().element(0)
          group.add_element().CopyFrom(root)
        db_key = db.Key._FromPb(key)
        update_entity = Update(
            key_name=str(db_key), parent=db_key)._populate_entity()
        extra_entities.append(update_entity._ToPb())
      request.entity_list().extend(extra_entities)
    elif call == "Delete" and not cls.ignore_deletes:
      cls.ignore_deletes = True
      try:
        if request.has_transaction():
          for key in request.key_list():
            update = Update(
                key_name=str(db.Key._FromPb(key)),
                parent=db.Key._FromPb(key))
            db.put(update)
        else:
          keys = request.key_list()
          request.clear_key()
          for key in keys:
            # TODO: group keys by entity groups
            def tx():
              db_key = db.Key._FromPb(key)
              update = Update(key_name=str(db_key),  parent=db_key)
              db.put(update)
              db.delete(db_key)
            db.run_in_transaction(tx)
      finally:
        cls.ignore_deletes = False

  @classmethod
  def post_hook(cls, service, call, request, response):
    if call == "Put" and not cls.ignore_deletes:
      keys = response.key_list()
      response.clear_key()
      keys = ([key for key in keys
        if key.path().element_list()[-1].type() != Update.kind()])
      response.key_list().extend(keys)

  @classmethod
  def init(cls):
    apiproxy_stub_map.apiproxy.GetPreCallHooks().Append(
        'trigger_manager', cls.pre_hook, 'datastore_v3')
    apiproxy_stub_map.apiproxy.GetPostCallHooks().Append(
        'trigger_manager', cls.post_hook, 'datastore_v3')

  @classmethod
  def init_if_needed(cls):
    if not cls.initialized:
      cls.initialized = True
      cls.init()

  @classmethod
  def register_trigger(cls, entity_kind, function_name):
    cls.init_if_needed()
    cls.triggers.setdefault(entity_kind, set())
    cls.triggers[entity_kind].add(function_name)

  @classmethod
  def get_triggers(cls, entity_kind):
    return cls.triggers.get(entity_kind)

  @classmethod
  def process_updates(cls):
    cls.ignore_deletes = True
    try:
      while True:
        # TODO time limits are needed here
        updates = Update.all().fetch(limit=50)
        if not updates:
          return
        for update in updates:
          def tx():
            db.delete(update.key())
            taskqueue.add(url=UPDATE_URL, params={'key': update.key().name()},
                transactional=True)
          db.run_in_transaction(tx)
    finally:
      cls.ignore_deletes = False


class ProcessUpdatesHandler(webapp.RequestHandler):
  def get(self):
    TriggerManager.process_updates()


class UpdateEntityHandler(webapp.RequestHandler):
  def post(self):
    key = db.Key(self.request.get("key"))
    triggers = TriggerManager.get_triggers(key.kind())
    if not triggers:
      return
    for trigger in triggers:
      util.for_name(trigger)(key)


def trigger(kind):
  """Trigger registering decorator.

  Args:
    kind: entity kind to attach defined trigger to.
  """
  TriggerManager.init_if_needed()
  def decorator(f):
    trigger_name = util.fq_name(f)
    TriggerManager.register_trigger(kind, trigger_name)
    return f
  return decorator


application = webapp.WSGIApplication([
  (PROCESS_URL, ProcessUpdatesHandler),
  (UPDATE_URL, UpdateEntityHandler),
  ], debug=True)
