[beam] branch master updated: [BEAM-12445] Handle none bq_client (#15194)

[goenka]

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

goenka 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 46842eb  [BEAM-12445] Handle none bq_client (#15194)
46842eb is described below

commit 46842eb2878021962e619d0bc863c31da6dd474b
Author: Ankur <angoe...@users.noreply.github.com>
AuthorDate: Wed Jul 21 00:56:47 2021 -0700

    [BEAM-12445] Handle none bq_client (#15194)
    
    * [BEAM-12445] Handle none bq_client
    
    * fix test case
---
 sdks/python/apache_beam/io/gcp/bigquery_tools.py      | 13 ++++++++-----
 sdks/python/apache_beam/io/gcp/bigquery_tools_test.py |  4 +---
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py 
b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index c9540ce..eefdb05 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -641,14 +641,17 @@ class BigQueryWrapper(object):
 
     started_millis = int(time.time() * 1000)
     try:
-      table_ref = gcp_bigquery.DatasetReference(project_id,
-                                                dataset_id).table(table_id)
+      table_ref_str = '%s.%s.%s' % (project_id, dataset_id, table_id)
       errors = self.gcp_bq_client.insert_rows_json(
-          table_ref, json_rows=rows, row_ids=insert_ids, 
skip_invalid_rows=True)
+          table_ref_str,
+          json_rows=rows,
+          row_ids=insert_ids,
+          skip_invalid_rows=True)
       if not errors:
         service_call_metric.call('ok')
-      for insert_error in errors:
-        service_call_metric.call(insert_error['errors'][0])
+      else:
+        for insert_error in errors:
+          service_call_metric.call(insert_error['errors'][0])
     except HttpError as e:
       service_call_metric.call(e)
 
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py 
b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
index 566e2b3..c9ed827 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -51,7 +51,6 @@ from apache_beam.options.value_provider import 
StaticValueProvider
 # pylint: disable=wrong-import-order, wrong-import-position
 try:
   from apitools.base.py.exceptions import HttpError, HttpForbiddenError
-  from google.cloud import bigquery as gcp_bigquery
 except ImportError:
   HttpError = None
   HttpForbiddenError = None
@@ -860,8 +859,7 @@ class TestBigQueryWriter(unittest.TestCase):
 
     sample_row = {'i': 1, 'b': True, 's': 'abc', 'f': 3.14}
     client.insert_rows_json.assert_called_with(
-        gcp_bigquery.TableReference(
-            gcp_bigquery.DatasetReference('project', 'dataset'), 'table'),
+        '%s.%s.%s' % ('project', 'dataset', 'table'),
         json_rows=[sample_row],
         row_ids=['_1'],
         skip_invalid_rows=True)