payload_utils/config_to_datastore.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright 2021 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Marshal ConfigBundles and store them into the UFS datastore through the Google datastore API.

By default, this function reads a ConfigBundleList from
'hw_design/generated/configs.jsonproto' and parses ConfigBundles from it. Using
the API specified, it encodes the ConfigBundle as a ConfigBundleEntity and
stores it into the UFS datastore.
"""

import argparse
import datetime
import logging

from google.cloud import datastore

from checker import io_utils
from common import proto_utils

# type constants
INPUT_TYPE = 'chromiumos.config.payload.ConfigBundleList'
OUTPUT_TYPE = 'chromiumos.config.payload.ConfigBundle'

# UFS services
UFS_DEV_PROJECT = 'unified-fleet-system-dev'
UFS_PROD_PROJECT = 'unified-fleet-system'

# datastore constants
CONFIG_BUNDLE_KIND = 'ConfigBundle'


def get_ufs_project(env):
  """Return project name based on env argument."""
  if env == 'dev':
    return UFS_DEV_PROJECT
  if env == 'prod':
    return UFS_PROD_PROJECT
  raise RuntimeError('get_ufs_project: environment %s not supported' % env)


def generate_entity_id(bundle):
  """Generate ConfigBundleEntity id as ${program_id}-${design_id}."""
  return bundle.design_list[0].program_id.value + '-' + bundle.design_list[
      0].id.value


def update_config_bundle(bundle, project):
  """Take a ConfigBundle and store it in the UFS datastore as a ConfigBundleEntity."""
  eid = generate_entity_id(bundle)
  logging.info('update_config_bundle: handling %s', eid)

  client = datastore.Client(project=project,)
  key = client.key(CONFIG_BUNDLE_KIND, eid)
  entity = datastore.Entity(
      key=key,
      exclude_from_indexes=['ConfigData'],
  )
  entity['ConfigData'] = bundle.SerializeToString()
  entity['updated'] = datetime.datetime.now()

  try:
    logging.info('update_config_bundle: putting entity into datastore for %s',
                 eid)
    client.put(entity)
  except Exception as err:  # pylint: disable=broad-except
    logging.exception('update_config_bundle: %s', err)


if __name__ == '__main__':
  logging.basicConfig(level=logging.INFO)
  parser = argparse.ArgumentParser(
      description=__doc__,
      formatter_class=argparse.RawDescriptionHelpFormatter,
  )

  parser.add_argument(
      'input',
      type=str,
      help='message to read in jsonproto format',
  )

  parser.add_argument(
      '--env',
      type=str,
      default='dev',
      help='environment flag for UFS service',
  )

  # load database of protobuffer name -> Type
  protodb = proto_utils.create_symbol_db()

  options = parser.parse_args()
  output = protodb.GetSymbol(OUTPUT_TYPE)()

  # read ConfigBundleList from options.input
  cb_list = io_utils.read_json_proto(
      protodb.GetSymbol(INPUT_TYPE)(), options.input)

  for config_bundle in cb_list.values:
    update_config_bundle(config_bundle, get_ufs_project(options.env))