Skip to content

Commit

Permalink
[BEAM-12445] Handle none bq_client (#15194)
Browse files Browse the repository at this point in the history
* [BEAM-12445] Handle none bq_client

* fix test case
  • Loading branch information
angoenka authored Jul 21, 2021
1 parent 6fcd339 commit 46842eb
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 8 deletions.
13 changes: 8 additions & 5 deletions sdks/python/apache_beam/io/gcp/bigquery_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -641,14 +641,17 @@ def _insert_all_rows(

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)

Expand Down
4 changes: 1 addition & 3 deletions sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,6 @@
# 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
Expand Down Expand Up @@ -860,8 +859,7 @@ def test_rows_are_written(self):

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)
Expand Down

0 comments on commit 46842eb

Please sign in to comment.