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)