boxster: update script to parse DutAttributes into UFS datastore
Added capability to parse and extract DutAttributes from DutAttributeList
into UFS datastore. IDs are based on the DutAttribute ID value.
BUG=b/186663540
TEST=manual testing
Change-Id: I6f40efb4c700c2ee48ea19f3ac4e72591b348c63
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/config/+/3043623
Commit-Queue: Justin Suen <justinsuen@google.com>
Reviewed-by: Sean McAllister <smcallis@google.com>
diff --git a/payload_utils/config_to_datastore.py b/payload_utils/config_to_datastore.py
index 9b51b9f..958c1f0 100755
--- a/payload_utils/config_to_datastore.py
+++ b/payload_utils/config_to_datastore.py
@@ -4,12 +4,18 @@
# 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.
+"""Marshal ConfigBundles and DutAttributes. 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.
+By default, this script converts a few config-related protos into datastore
+entities:
+
+1. ConfigBundleList from 'hw_design/generated/configs.jsonproto'
+2. DutAttributeList from 'dut_attributes/generated/dut_attributes.jsonproto'
+
+The lists are parsed and individual entities are extracted. Using the datastore
+client specified, it encodes the protos as datastore entities and stores them
+into the UFS datastore.
"""
import argparse
@@ -22,8 +28,10 @@
from common import proto_utils
# type constants
-INPUT_TYPE = 'chromiumos.config.payload.ConfigBundleList'
-OUTPUT_TYPE = 'chromiumos.config.payload.ConfigBundle'
+CB_INPUT_TYPE = 'chromiumos.config.payload.ConfigBundleList'
+CB_OUTPUT_TYPE = 'chromiumos.config.payload.ConfigBundle'
+DA_INPUT_TYPE = 'chromiumos.test.api.DutAttributeList'
+DA_OUTPUT_TYPE = 'chromiumos.test.api.DutAttribute'
# UFS services
UFS_DEV_PROJECT = 'unified-fleet-system-dev'
@@ -31,6 +39,7 @@
# datastore constants
CONFIG_BUNDLE_KIND = 'ConfigBundle'
+DUT_ATTRIBUTE_KIND = 'DutAttribute'
def get_ufs_project(env):
@@ -48,6 +57,17 @@
0].id.value
+def handle_config_bundle_list(cb_list_path, env):
+ """Take a path to a ConfigBundleList, iterate through the list and store into
+ UFS datastore based on env.
+ """
+ cb_list = io_utils.read_json_proto(
+ protodb.GetSymbol(CB_INPUT_TYPE)(), cb_list_path)
+
+ for config_bundle in cb_list.values:
+ update_config_bundle(config_bundle, get_ufs_project(env))
+
+
def update_config_bundle(bundle, project):
"""Take a ConfigBundle and store it in the UFS datastore as a ConfigBundleEntity."""
eid = generate_entity_id(bundle)
@@ -62,12 +82,39 @@
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)
+ logging.info('update_config_bundle: putting entity into datastore for %s',
+ eid)
+ client.put(entity)
+
+
+def handle_dut_attribute_list(dut_attr_list_path, env):
+ """Take a path to a DutAttributeList, iterate through the list and store into
+ UFS datastore based on env.
+ """
+ dut_attr_list = io_utils.read_json_proto(
+ protodb.GetSymbol(DA_INPUT_TYPE)(), dut_attr_list_path)
+
+ for dut_attribute in dut_attr_list.dut_attributes:
+ update_dut_attribute(dut_attribute, get_ufs_project(env))
+
+
+def update_dut_attribute(attr, project):
+ """Take a DutAttribute and store it in the UFS datastore as a DutAttributeEntity."""
+ eid = attr.id.value
+ logging.info('update_dut_attribute: handling %s', eid)
+
+ client = datastore.Client(project=project,)
+ key = client.key(DUT_ATTRIBUTE_KIND, eid)
+ entity = datastore.Entity(
+ key=key,
+ exclude_from_indexes=['AttributeData'],
+ )
+ entity['AttributeData'] = attr.SerializeToString()
+ entity['updated'] = datetime.datetime.now()
+
+ logging.info('update_dut_attribute: putting entity into datastore for %s',
+ eid)
+ client.put(entity)
if __name__ == '__main__':
@@ -78,12 +125,6 @@
)
parser.add_argument(
- 'input',
- type=str,
- help='message to read in jsonproto format',
- )
-
- parser.add_argument(
'--env',
type=str,
default='dev',
@@ -92,13 +133,9 @@
# 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))
+ handle_config_bundle_list("hw_design/generated/configs.jsonproto",
+ options.env)
+ handle_dut_attribute_list("dut_attributes/generated/dut_attributes.jsonproto",
+ options.env)