appengine/swarming/server/resultdb.py

# Copyright 2020 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.
"""Functions to communicate with ResultDB for swarming tasks."""

import logging
import uuid

from google.appengine.api import app_identity
from google.appengine.ext import ndb

from components import net
from components import prpc
from components import utils

from server import config


@ndb.tasklet
def create_invocation_async(task_run_id, realm, deadline):
  """This is wrapper for CreateInvocation API.

  Returns:
    update-token for created invocation.
  """
  hostname = app_identity.get_default_version_hostname()
  response_headers = {}

  assert deadline.utcoffset() is None, (
    "non-default timezone is not supported, %s" % deadline.utcoffset()
  )
  yield _call_resultdb_recorder_api_async(
    "CreateInvocation",
    {
      "requestId": str(uuid.uuid4()),
      "invocationId": _get_invocation_id(task_run_id),
      "invocation": {
        "producerResource": "//%s/tasks/%s" % (hostname, task_run_id),
        "realm": realm,
        "deadline": deadline.isoformat() + "Z",
      },
    },
    project_id=realm.split(":")[0],
    response_headers=response_headers,
  )
  update_token = response_headers.get("update-token")
  assert update_token, (
    "response_headers should have valid update-token: %s" % response_headers
  )
  raise ndb.Return(update_token)


@ndb.tasklet
def finalize_invocation_async(task_run_id, update_token):
  """This is wrapper for FinalizeInvocation API."""

  try:
    invocation_name = get_invocation_name(task_run_id)
    yield _call_resultdb_recorder_api_async(
      "FinalizeInvocation",
      {
        "name": invocation_name,
      },
      headers={"update-token": update_token},
      scopes=None,
    )
  except net.Error:
    logging.exception("Failed to finalize %s", invocation_name)


def get_invocation_name(task_run_id):
  return "invocations/%s" % _get_invocation_id(task_run_id)


### Private code


def _get_invocation_id(task_run_id):
  hostname = app_identity.get_default_version_hostname()
  return "task-%s-%s" % (hostname, task_run_id)


@ndb.tasklet
def _call_resultdb_recorder_api_async(
  method,
  request,
  headers=None,
  response_headers=None,
  scopes=(net.EMAIL_SCOPE,),
  project_id=None,
):
  cfg = config.settings()
  rdb_url = cfg.resultdb.server
  if not rdb_url:
    raise ValueError("ResultDB integration is not configured")

  utils.validate_root_service_url(rdb_url)

  # See Recoder API for ResultDB in
  # https://chromium.googlesource.com/infra/luci/luci-go/+/HEAD/resultdb/proto/rpc/v1/recorder.proto
  # But beware that proto JSON serialization uses camelCase, not snake_case.
  yield net.json_request_async(
    url="%s/prpc/luci.resultdb.v1.Recorder/%s" % (rdb_url, method),
    method="POST",
    payload=request,
    scopes=scopes,
    headers=headers,
    response_headers=response_headers,
    project_id=project_id,
  )