def test_rows_are_written(self):
client = mock.Mock()
table = bigquery.Table(tableReference=bigquery.TableReference(
projectId='project', datasetId='dataset', tableId='table'),
schema=bigquery.TableSchema())
client.tables.Get.return_value = table
write_disposition = beam.io.BigQueryDisposition.WRITE_APPEND
insert_response = mock.Mock()
insert_response.insertErrors = []
client.tabledata.InsertAll.return_value = insert_response
with beam.io.BigQuerySink(
'project:dataset.table',
write_disposition=write_disposition).writer(client) as writer:
writer.Write({'i': 1, 'b': True, 's': 'abc', 'f': 3.14})
sample_row = {'i': 1, 'b': True, 's': 'abc', 'f': 3.14}
expected_rows = []
json_object = bigquery.JsonObject()
for k, v in iteritems(sample_row):
json_object.additionalProperties.append(
bigquery.JsonObject.AdditionalProperty(key=k,
value=to_json_value(v)))
expected_rows.append(
bigquery.TableDataInsertAllRequest.RowsValueListEntry(
insertId='_1', # First row ID generated with prefix ''
json=json_object))
client.tabledata.InsertAll.assert_called_with(
bigquery.BigqueryTabledataInsertAllRequest(
projectId='project',
datasetId='dataset',
tableId='table',
tableDataInsertAllRequest=bigquery.TableDataInsertAllRequest(
rows=expected_rows)))
def _insert_all_rows(self, project_id, dataset_id, table_id, rows):
# The rows argument is a list of
# bigquery.TableDataInsertAllRequest.RowsValueListEntry instances as
# required by the InsertAll() method.
request = bigquery.BigqueryTabledataInsertAllRequest(
projectId=project_id, datasetId=dataset_id, tableId=table_id,
tableDataInsertAllRequest=bigquery.TableDataInsertAllRequest(
# TODO(silviuc): Should have an option for skipInvalidRows?
# TODO(silviuc): Should have an option for ignoreUnknownValues?
rows=rows))
response = self.client.tabledata.InsertAll(request)
# response.insertErrors is not [] if errors encountered.
return not response.insertErrors, response.insertErrors
def _insert_all_rows(self, project_id, dataset_id, table_id, rows,
skip_invalid_rows=False):
"""Calls the insertAll BigQuery API endpoint.
Docs for this BQ call: https://cloud.google.com/bigquery/docs/reference\
/rest/v2/tabledata/insertAll."""
# The rows argument is a list of
# bigquery.TableDataInsertAllRequest.RowsValueListEntry instances as
# required by the InsertAll() method.
request = bigquery.BigqueryTabledataInsertAllRequest(
projectId=project_id, datasetId=dataset_id, tableId=table_id,
tableDataInsertAllRequest=bigquery.TableDataInsertAllRequest(
skipInvalidRows=skip_invalid_rows,
# TODO(silviuc): Should have an option for ignoreUnknownValues?
rows=rows))
response = self.client.tabledata.InsertAll(request)
# response.insertErrors is not [] if errors encountered.
return not response.insertErrors, response.insertErrors
