ariofrio/oration

"""
Sets up a standard set of routes according to the Cicero API.

Posts tasks to a push Task Queue, the data is stored in the Data Store.

(Automatically generated by Oration.)
"""

import os
import sys

import logging
import base64
from datetime import datetime

import json
import wsgiref.handlers
import webapp2

from google.appengine.api import taskqueue
from google.appengine.ext import db

import {{ namespace }}

class TaskInfo(db.Model):
  status = db.StringProperty()
  start_time = db.StringProperty()
  finish_time = db.StringProperty()

class Text(db.Model):
  content = db.TextProperty()

class TaskRoute(webapp2.RequestHandler):
  def get(self):
    id = self.request.get('id') or self.request.get('task_id')
    logging.info("Getting status of task " + id)
    model = TaskInfo.get_by_key_name(id)

    result = None
    if model:
      task_info = model.__dict__['_entity']
      task_info['result'] = 'success'
      result = task_info
    else:
      result = {'result':'failure', 'reason':'not found'}
    self.response.out.write(json.dumps(result))

  def put(self):
    allowed_functions = ['{{ function_name }}']
    function = self.request.get('function') or self.request.get('f')

    key_length = 16
    id = self.request.get('id')
    if not id:
      id = base64.b64encode(os.urandom(key_length), '-_')
    output_location = self.request.get('output_location') or self.request.get('output')
    if not output_location:
      output_location = base64.b64encode(os.urandom(key_length), '-_')

    logging.info("Posting task " + id + " for function " + function)

    if function in allowed_functions:
      taskqueue.add(url='/compute', params={'data': json.dumps(
        {'function': function, 'id': id,
         'output_location': output_location})})

      result = {'result': 'success', 'id': id,
                'output_location': output_location}
      self.response.out.write(json.dumps(result))
    else:
      reason = 'Cannot add a task for function type ' + str(function)
      result ={'result': 'failure', 'reason': reason}
      self.response.out.write(json.dumps(result))

  def post(self):
    self.put()

  def delete(self):
    # REGRESSION:
    logging.info("NOT cancelling task: not implemented")
    result = {'result': 'not implemented'}
    self.response.out.write(result)


class DataRoute(webapp2.RequestHandler):
  def get(self):
    location = self.request.get('location')
    logging.info("Getting data at location " + location)
    model = Text.get_by_key_name(location)

    result = None
    if model:
      result = {'result': 'success', 'output': model.content}
    else:
      result = {'result': 'failure', 'reason': 'not found'}
    self.response.out.write(json.dumps(result))

  def put(self):
    location = self.request.get('location')
    logging.info("Putting data into location " + location)
    model = Text(key_name = location)
    model.content = self.request.get('content')
    model.put()
    self.response.out.write(json.dumps({'result': 'success'}))

  def delete(self):
    location = self.request.get('location')
    logging.info("Deleting data at location " + location)
    model = Text(key_name = location)

    result = None
    if model:
      model.delete()
      result = {'result': 'success'}
    else:
      result = {'result': 'failure', 'reason': 'not found'}
    self.response.out.write(result)

class ComputeRoute(webapp2.RequestHandler):
  def post(self):
    data = json.loads(self.request.get('data'))
    id = data['id']
    function = data['function']
    output_location = data['output_location']
    logging.info("Starting task " + id + " for function " + function)

    logging.info("Updating status of task " + id + ": started")
    TaskInfo(key_name = id, status = 'started', start_time = datetime.now().isoformat()).put()

    logging.info("Actually running task " + id)
    content = str({{ namespace }}.{{ function_name }}())

    logging.info("Putting output data into location " + output_location)
    Text(key_name = output_location, content = content).put()

    logging.info("Updating status of task " + id + ": finished")
    status = TaskInfo.get_by_key_name(id)
    status.status = 'finished'
    status.finish_time = datetime.now().isoformat()
    status.put()


class IndexPage(webapp2.RequestHandler):
  def get(self):
    self.response.out.write("Hello!")

logging.getLogger().setLevel(logging.INFO)
app = webapp2.WSGIApplication([('/task', TaskRoute),
                              ('/data', DataRoute),
                              ('/compute', ComputeRoute),
                              ('/', IndexPage),
                              ], debug=True)
def main():
  app.run()

if __name__ == '__main__':
  main()