[BEAM-12670] Fix regression, properly report API call metric properly on error for BQ streaming_inserts

Alex Amato · Alex Amato · commit 72dbbf2e3ce9 · 2021-07-28T10:27:57.000-07:00
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -617,6 +617,8 @@ def _insert_all_rows(
 
     Docs for this BQ call: https://cloud.google.com/bigquery/docs/reference\
       /rest/v2/tabledata/insertAll."""
+    from google.api_core.exceptions import ClientError
+    from google.api_core.exceptions import GoogleAPICallError
     # The rows argument is a list of
     # bigquery.TableDataInsertAllRequest.RowsValueListEntry instances as
     # required by the InsertAll() method.
@@ -652,9 +654,13 @@ def _insert_all_rows(
       else:
         for insert_error in errors:
           service_call_metric.call(insert_error['errors'][0])
+    except (ClientError, GoogleAPICallError) as e:
+      # e.code.value contains the numeric http status code.
+      service_call_metric.call(e.code.value)
+      # Re-reise the exception so that we re-try appropriately.
+      raise
     except HttpError as e:
       service_call_metric.call(e)
-
       # Re-reise the exception so that we re-try appropriately.
       raise
     finally:
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -33,6 +33,7 @@
 
 import apache_beam as beam
 from apache_beam.internal.gcp.json_value import to_json_value
+from apache_beam.io.gcp import resource_identifiers
 from apache_beam.io.gcp.bigquery import TableRowJsonCoder
 from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR
 from apache_beam.io.gcp.bigquery_tools import AvroRowWriter
@@ -44,6 +45,8 @@
 from apache_beam.io.gcp.bigquery_tools import parse_table_reference
 from apache_beam.io.gcp.bigquery_tools import parse_table_schema_from_json
 from apache_beam.io.gcp.internal.clients import bigquery
+from apache_beam.metrics import monitoring_infos
+from apache_beam.metrics.execution import MetricsEnvironment
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.value_provider import StaticValueProvider
 
@@ -422,6 +425,62 @@ def test_perform_load_job_with_source_stream(self):
     upload = client.jobs.Insert.call_args[1]["upload"]
     self.assertEqual(b'some,data', upload.stream.read())
 
+  def verify_write_call_metric(
+      self, project_id, dataset_id, table_id, status, count):
+    """Check if an metric was recorded for the BQ IO write API call."""
+    process_wide_monitoring_infos = list(
+        MetricsEnvironment.process_wide_container().
+        to_runner_api_monitoring_infos(None).values())
+    resource = resource_identifiers.BigQueryTable(
+        project_id, dataset_id, table_id)
+    labels = {
+        # TODO(ajamato): Add Ptransform label.
+        monitoring_infos.SERVICE_LABEL: 'BigQuery',
+        # Refer to any method which writes elements to BigQuery in batches
+        # as "BigQueryBatchWrite". I.e. storage API's insertAll, or future
+        # APIs introduced.
+        monitoring_infos.METHOD_LABEL: 'BigQueryBatchWrite',
+        monitoring_infos.RESOURCE_LABEL: resource,
+        monitoring_infos.BIGQUERY_PROJECT_ID_LABEL: project_id,
+        monitoring_infos.BIGQUERY_DATASET_LABEL: dataset_id,
+        monitoring_infos.BIGQUERY_TABLE_LABEL: table_id,
+        monitoring_infos.STATUS_LABEL: status,
+    }
+    expected_mi = monitoring_infos.int64_counter(
+        monitoring_infos.API_REQUEST_COUNT_URN, count, labels=labels)
+    expected_mi.ClearField("start_time")
+
+    found = False
+    for actual_mi in process_wide_monitoring_infos:
+      actual_mi.ClearField("start_time")
+      if expected_mi == actual_mi:
+        found = True
+        break
+    self.assertTrue(
+        found, "Did not find write call metric with status: %s" % status)
+
+  def test_insert_rows_sets_metric_on_failure(self):
+    from google.api_core import exceptions
+    MetricsEnvironment.process_wide_container().reset()
+    client = mock.Mock()
+    client.insert_rows_json = mock.Mock(
+        # Fail a few times, then succeed.
+        side_effect=[
+            exceptions.DeadlineExceeded("Deadline Exceeded"),
+            exceptions.InternalServerError("Internal Error"),
+            [],
+        ])
+    wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client)
+    wrapper.insert_rows("my_project", "my_dataset", "my_table", [])
+
+    # Expect two failing calls, then a success (i.e. two retries).
+    self.verify_write_call_metric(
+        "my_project", "my_dataset", "my_table", "deadline_exceeded", 1)
+    self.verify_write_call_metric(
+        "my_project", "my_dataset", "my_table", "internal", 1)
+    self.verify_write_call_metric(
+        "my_project", "my_dataset", "my_table", "ok", 1)
+
 
 @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
 class TestBigQueryReader(unittest.TestCase):
