server.py

#!/usr/bin/env python3
"""Web server for the export-to-Google-Drive EE demo application.

The code in this file runs on App Engine. It's called when the user loads the
web page, requests lights from a different year, or requests an export.

Our App Engine code does most of the communication with EE. It uses the
EE Python library and the service account specified in config.py. The
exception is that when the browser loads map tiles it talks directly with EE.

The app uses two different sets of credentials:

 1) The service account credentials, which are used to query the EE API and
    share exported files in Google Drive.
 2) The end user's OAuth2 credentials for Google Drive. These are used to copy
    exported files from the service account's Drive into the user's.

Common flows:

When the user first loads the webpage, they will be asked to authorize the
app's access to Google Drive (yielding credentials #2 above). The decorator
@OAUTH_DECORATOR.oauth_required triggers this flow. The main handler then
generates a unique client ID for the Channel API connection, injects it
into the index.html template, and returns the page contents.

When the user changes the year in the UI, a map ID is generated by the
/mapid handler.

When the user exports a file, the /export handler then kicks off an export
runner (running asynchronously) to create the EE task and poll for the task's
completion. When the EE task completes, the file is copied from the service
account's Drive folder to the user's Drive folder and an update is sent to the
user's browser using the Channel API.
"""

import json
import logging
import os
import random
import string
import time

import config
import drive
import ee
import jinja2
import oauth2client.appengine
import webapp2

from google.appengine.api import channel
from google.appengine.api import taskqueue
from google.appengine.api import urlfetch
from google.appengine.api import users


###############################################################################
#                               Initialization.                               #
###############################################################################


# The URL fetch timeout (seconds).
URL_FETCH_TIMEOUT = 60

# Our App Engine service account's credentials for use with Earth Engine.
EE_CREDENTIALS = ee.ServiceAccountCredentials(
    config.EE_ACCOUNT, config.EE_PRIVATE_KEY_FILE)

# Initialize the EE API.
ee.Initialize(EE_CREDENTIALS)

# The Jinja templating system we use to dynamically generate HTML. See:
# http://jinja.pocoo.org/docs/dev/
JINJA2_ENVIRONMENT = jinja2.Environment(
    loader=jinja2.FileSystemLoader(os.path.dirname(__file__)), autoescape=True)

# Check https://developers.google.com/drive/scopes for all available scopes.
OAUTH_SCOPE = 'https://www.googleapis.com/auth/drive'

# The app's service account credentials (for Google Drive).
APP_CREDENTIALS = oauth2client.client.SignedJwtAssertionCredentials(
    config.EE_ACCOUNT,
    open(config.EE_PRIVATE_KEY_FILE, 'rb').read(),
    OAUTH_SCOPE)

# An authenticated Drive helper object for the app service account.
APP_DRIVE_HELPER = drive.DriveHelper(APP_CREDENTIALS)

# The decorator to trigger the user's Drive permissions request flow.
OAUTH_DECORATOR = oauth2client.appengine.OAuth2Decorator(
    client_id=config.OAUTH_CLIENT_ID,
    client_secret=config.OAUTH_CLIENT_SECRET,
    scope=OAUTH_SCOPE)

# The ImageCollection of night-time lights images.
IMAGE_COLLECTION_ID = 'NOAA/DMSP-OLS/NIGHTTIME_LIGHTS'

# The resolution of the exported images (meters per pixel).
EXPORT_RESOLUTION = 30

# The maximum number of pixels in an exported image.
EXPORT_MAX_PIXELS = 1e10

# The visualization parameters for the images.
VIZ_PARAMS = {
    'min': 0,
    'max': 63,
}

# The frequency to poll for export EE task completion (seconds).
TASK_POLL_FREQUENCY = 10

# The image IDs within NOAA/DMSP-OLS/NIGHTTIME_LIGHTS, which are formatted
# slightly inconsistently.
IMAGE_IDS = [
    'F101992', 'F101993', 'F121994', 'F121995', 'F121996', 'F141997',
    'F141998', 'F141999', 'F152000', 'F152001', 'F152002', 'F152003',
    'F162004', 'F162005', 'F162006', 'F162007', 'F162008', 'F162009',
    'F182010', 'F182011', 'F182012'
]

###############################################################################
#                             Web request handlers.                           #
###############################################################################


class MainHandler(webapp2.RequestHandler):
  """A servlet to handle requests to load the main web page."""

  @OAUTH_DECORATOR.oauth_required
  def get(self):
    """Returns the main web page with Channel API details included."""
    client_id = _GetUniqueString()

    template = JINJA2_ENVIRONMENT.get_template('index.html')
    self.response.out.write(template.render({
        'channelToken': channel.create_channel(client_id),
        'clientId': client_id,
    }))


class DataHandler(webapp2.RequestHandler):
  """A servlet base class for responding to data queries.

  We use this base class to wrap our web request handlers with try/except
  blocks and set per-thread values (e.g. URL_FETCH_TIMEOUT).
  """

  def get(self):
    self.Handle(self.DoGet)

  def post(self):
    self.Handle(self.DoPost)

  def DoGet(self):
    """Processes a GET request and returns a JSON-encodable result."""
    raise NotImplementedError()

  def DoPost(self):
    """Processes a POST request and returns a JSON-encodable result."""
    raise NotImplementedError()

  @OAUTH_DECORATOR.oauth_required
  def Handle(self, handle_function):
    """Responds with the result of the handle_function or errors, if any."""
    # Note: The fetch timeout is thread-local so must be set separately
    # for each incoming request.
    urlfetch.set_default_fetch_deadline(URL_FETCH_TIMEOUT)
    try:
      response = handle_function()
    except Exception as e:  # pylint: disable=broad-except
      response = {'error': str(e)}
    if response:
      self.response.headers['Content-Type'] = 'application/json'
      self.response.out.write(json.dumps(response))


class MapIdHandler(DataHandler):
  """A servlet to handle requests for lights map IDs for a given year."""

  def DoGet(self):
    """Returns the map ID of an image for the requested year.

    HTTP Parameters:
      year: The year of the image, in "YYYY" format.

    Returns:
      A dictionary with two keys: mapid and token.
    """
    image = _GetImage(self.request.get('year'))
    mapid = image.getMapId(VIZ_PARAMS)
    return {'mapid': mapid['mapid'], 'token': mapid['token']}


class ExportHandler(DataHandler):
  """A servlet to handle requests for image exports."""

  def DoPost(self):
    """Kicks off export of an image for the specified year and region.

    HTTP Parameters:
      coordinates: The coordinates of the polygon to export.
      filename: The final filename of the file to create in the user's Drive.
      client_id: The ID of the client (for the Channel API).
      year: The year of the image, in "YYYY" format.
    """

    # Kick off an export runner to start and monitor the EE export task.
    # Note: The work "task" is used by both Earth Engine and App Engine to refer
    # to two different things. "TaskQueue" is an async App Engine service.
    taskqueue.add(url='/exportrunner', params={
        'coordinates': self.request.get('coordinates'),
        'filename': self.request.get('filename'),
        'client_id': self.request.get('client_id'),
        'year': self.request.get('year'),
        'email': users.get_current_user().email(),
        'user_id': users.get_current_user().user_id(),
    })


###############################################################################
#                           The task status poller.                           #
###############################################################################


class ExportRunnerHandler(webapp2.RequestHandler):
  """A servlet for handling async export task requests."""

  def post(self):
    """Exports an image for the year and region, gives it to the user.

    This is called by our trusted export handler and runs as a separate
    process.

    HTTP Parameters:
      email: The email address of the user who initiated this task.
      filename: The final filename of the file to create in the user's Drive.
      client_id: The ID of the client (for the Channel API).
      task: The pickled task to poll.
      temp_file_prefix: The prefix of the temp file in the service account's
          Drive.
      user_id: The ID of the user who initiated this task.
    """
    coordinates = self.request.get('coordinates')
    filename = self.request.get('filename')
    client_id = self.request.get('client_id')
    email = self.request.get('email')
    user_id = self.request.get('user_id')
    year = self.request.get('year')

    # Get the image for the year and region to export.
    image = GetExportableImage(_GetImage(year), coordinates)

    # Use a unique prefix to identify the exported file.
    temp_file_prefix = _GetUniqueString()

    # Create and start the task.
    task = ee.batch.Export.image(
        image=image,
        description='Earth Engine Demo Export',
        config={
            'driveFileNamePrefix': temp_file_prefix,
            'maxPixels': EXPORT_MAX_PIXELS,
            'scale': EXPORT_RESOLUTION,
        })
    task.start()
    logging.info('Started EE task (id: %s).', task.id)

    # Wait for the task to complete (taskqueue auto times out after 10 mins).
    while task.active():
      logging.info('Polling for task (id: %s).', task.id)
      time.sleep(TASK_POLL_FREQUENCY)

    def _SendMessage(message):
      logging.info('Sent to client: ' + json.dumps(message))
      _SendMessageToClient(client_id, filename, message)

    # Make a copy (or copies) in the user's Drive if the task succeeded.
    state = task.status()['state']
    if state == ee.batch.Task.State.COMPLETED:
      logging.info('Task succeeded (id: %s).', task.id)
      try:
        link = _GiveFilesToUser(temp_file_prefix, email, user_id, filename)
        # Notify the user's browser that the export is complete.
        _SendMessage({'link': link})
      except Exception as e:  # pylint: disable=broad-except
        _SendMessage({'error': 'Failed to give file to user: ' + str(e)})
    else:
      _SendMessage({'error': 'Task failed (id: %s).' % task.id})


###############################################################################
#                               Routing table.                                #
###############################################################################


# The webapp2 routing table from URL paths to web request handlers. See:
# http://webapp-improved.appspot.com/tutorials/quickstart.html
app = webapp2.WSGIApplication([
    ('/export', ExportHandler),
    ('/exportrunner', ExportRunnerHandler),
    ('/mapid', MapIdHandler),
    ('/', MainHandler),
    (OAUTH_DECORATOR.callback_path, OAUTH_DECORATOR.callback_handler()),
])


###############################################################################
#                                   Helpers.                                  #
###############################################################################


def _GetImage(year):
  """Returns the night-time lights image for a given year.

  Args:
    year: The year for which to retrieve an image.

  Returns:
    An ee.Image with lights for the given year.
  """
  image_id = IMAGE_IDS[int(year) - 1992]
  return ee.Image(IMAGE_COLLECTION_ID + '/' + image_id).select(0)


def _GetUniqueString():
  """Returns a likely-to-be unique string."""
  random_str = ''.join(
      random.choice(string.ascii_uppercase + string.digits) for _ in range(6))
  date_str = str(int(time.time()))
  return date_str + random_str


def _SendMessageToClient(client_id, filename, params):
  """Sends a message to the client using the Channel API.

  Args:
    client_id: The ID of the client to message.
    filename: The name of the exported file the message is about.
    params: The params to send in the message (as a Dictionary).
  """
  params['filename'] = filename
  channel.send_message(client_id, json.dumps(params))


def GetExportableImage(image, coordinates):
  """Crops and formats the image for export.

  Args:
    image: The image to make exportable.
    coordinates: The coordinates to crop the image to.

  Returns:
    The export-ready image.
  """

  # Determine the geometry based on the polygon's coordinates.
  coordinates = json.loads(coordinates)
  geometry = ee.Geometry.Polygon(coordinates)

  # Compute the image to export based on parameters.
  clipped_image = image.clip(geometry)
  return clipped_image.visualize(**VIZ_PARAMS)


def _GiveFilesToUser(temp_file_prefix, email, user_id, filename):
  """Moves the files with the prefix to the user's Drive folder.

  Copies and then deletes the source files from the app's Drive.

  Args:
    temp_file_prefix: The prefix of the temp files in the service account's
      Drive.
    email: The email address of the user to give the files to.
    user_id: The ID of the user to give the files to.
    filename: The name to give the files in the user's Drive.

  Returns:
    A link to the files in the user's Drive.
  """
  files = APP_DRIVE_HELPER.GetExportedFiles(temp_file_prefix)

  # Grant the user write access to the file(s) in the app service
  # account's Drive.
  for f in files:
    APP_DRIVE_HELPER.GrantAccess(f['id'], email)

  # Create a Drive helper to access the user's Google Drive.
  user_credentials = oauth2client.appengine.StorageByKeyName(
      oauth2client.appengine.CredentialsModel,
      user_id, 'credentials').get()
  user_drive_helper = drive.DriveHelper(user_credentials)

  # Copy the file(s) into the user's Drive.
  if len(files) == 1:
    file_id = files[0]['id']
    copied_file_id = user_drive_helper.CopyFile(file_id, filename)
    trailer = 'open?id=' + copied_file_id
  else:
    trailer = ''
    for f in files:
      # The titles of the files include the coordinates separated by a dash.
      coords = '-'.join(f['title'].split('-')[-2:])
      user_drive_helper.CopyFile(f['id'], filename + '-' + coords)

  # Delete the file from the service account's Drive.
  for f in files:
    APP_DRIVE_HELPER.DeleteFile(f['id'])

  return 'https://drive.google.com/' + trailer