(beam) branch master updated: Save Job IDs of dataflow load tests (#29693)

[an...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

anandinguva pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new f93423009fd Save Job IDs of dataflow load tests (#29693)
f93423009fd is described below

commit f93423009fd2d73a76e44332b09f9fe38d75cd2b
Author: Anand Inguva <34...@users.noreply.github.com>
AuthorDate: Mon Dec 11 23:17:49 2023 +0000

    Save Job IDs of dataflow load tests (#29693)
    
    * Publish job (#29685)
    
    * Update schema if the default schema is different than the table schema
    
    * Add job id label to the BQ publisher for dataflow jobs
    
    * Append job id to each metric
    
    * Fix lint
---
 .../testing/load_tests/load_test_metrics_utils.py  | 28 ++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/sdks/python/apache_beam/testing/load_tests/load_test_metrics_utils.py b/sdks/python/apache_beam/testing/load_tests/load_test_metrics_utils.py
index 1ff46a3f7d1..d1da4667dcb 100644
--- a/sdks/python/apache_beam/testing/load_tests/load_test_metrics_utils.py
+++ b/sdks/python/apache_beam/testing/load_tests/load_test_metrics_utils.py
@@ -33,6 +33,8 @@ import json
 import logging
 import time
 import uuid
+from typing import Any
+from typing import Dict
 from typing import List
 from typing import Mapping
 from typing import Optional
@@ -45,6 +47,7 @@ import apache_beam as beam
 from apache_beam.metrics import Metrics
 from apache_beam.metrics.metric import MetricResults
 from apache_beam.metrics.metric import MetricsFilter
+from apache_beam.runners.dataflow.dataflow_runner import DataflowPipelineResult
 from apache_beam.runners.runner import PipelineResult
 from apache_beam.transforms.window import TimestampedValue
 from apache_beam.utils.timestamp import Timestamp
@@ -65,6 +68,7 @@ ID_LABEL = 'test_id'
 SUBMIT_TIMESTAMP_LABEL = 'timestamp'
 METRICS_TYPE_LABEL = 'metric'
 VALUE_LABEL = 'value'
+JOB_ID_LABEL = 'job_id'
 
 SCHEMA = [{
     'name': ID_LABEL, 'field_type': 'STRING', 'mode': 'REQUIRED'
@@ -80,6 +84,8 @@ SCHEMA = [{
               'mode': 'REQUIRED'
           }, {
               'name': VALUE_LABEL, 'field_type': 'FLOAT', 'mode': 'REQUIRED'
+          }, {
+              'name': JOB_ID_LABEL, 'field_type': 'STRING', 'mode': 'NULLABLE'
           }]
 
 _LOGGER = logging.getLogger(__name__)
@@ -254,13 +260,27 @@ class MetricsReader(object):
     # Under each key there is list of objects of each metric type. It is
     # required to prepare metrics for publishing purposes. Expected is to have
     # a list of dictionaries matching the schema.
+
     insert_dicts = self._prepare_all_metrics(metrics, metric_id)
 
     insert_dicts += self._prepare_extra_metrics(metric_id, extra_metrics)
+
+    # Add job id for dataflow jobs for easier debugging.
+    job_id = None
+    if isinstance(result, DataflowPipelineResult):
+      job_id = result.job_id()
+      self._add_job_id_to_metrics(insert_dicts, job_id)
+
     if len(insert_dicts) > 0:
       for publisher in self.publishers:
         publisher.publish(insert_dicts)
 
+  def _add_job_id_to_metrics(self, metrics: List[Dict[str, Any]],
+                             job_id) -> List[Dict[str, Any]]:
+    for metric in metrics:
+      metric[JOB_ID_LABEL] = job_id
+    return metrics
+
   def _prepare_extra_metrics(
       self, metric_id: str, extra_metrics: Optional[dict] = None):
     ts = time.time()
@@ -479,6 +499,12 @@ class BigQueryClient(object):
       table = bigquery.Table(table_ref, schema=bq_schemas)
       self._bq_table = self._client.create_table(table)
 
+  def _update_schema(self):
+    table_schema = self._bq_table.schema
+    if self.schema and len(table_schema) != self.schema:
+      self._bq_table.schema = self._prepare_schema()
+      self._bq_table = self._client.update_table(self._bq_table, ["schema"])
+
   def _get_dataset(self, dataset_name):
     bq_dataset_ref = self._client.dataset(dataset_name)
     try:
@@ -490,6 +516,8 @@ class BigQueryClient(object):
     return bq_dataset
 
   def save(self, results):
+    # update schema if needed
+    self._update_schema()
     return self._client.insert_rows(self._bq_table, results)