Merge pull request #90 from hagino3000/add_stream_insert_options

Support ignoreUnknownValues and skipInvalidRows option for insertAll

Author: tylertreat

bigquery/client.py

@@ -1144,7 +1144,8 @@ def wait_for_job(self, job, interval=5, timeout=60):
 
         return job_resource
 
-    def push_rows(self, dataset, table, rows, insert_id_key=None):
+    def push_rows(self, dataset, table, rows, insert_id_key=None,
+                  skip_invalid_rows=None, ignore_unknown_values=None):
         """Upload rows to BigQuery table.
 
         Parameters
@@ -1157,6 +1158,10 @@ def push_rows(self, dataset, table, rows, insert_id_key=None):
             A ``list`` of rows (``dict`` objects) to add to the table
         insert_id_key : str, optional
             Key for insertId in row
+        skip_invalid_rows : bool, optional
+            Insert all valid rows of a request, even if invalid rows exist.
+        ignore_unknown_values : bool, optional
+            Accept rows that contain values that do not match the schema.
 
         Returns
         -------
@@ -1180,6 +1185,12 @@ def push_rows(self, dataset, table, rows, insert_id_key=None):
             "rows": rows_data
         }
 
+        if skip_invalid_rows is not None:
+            data['skipInvalidRows'] = skip_invalid_rows
+
+        if ignore_unknown_values is not None:
+            data['ignoreUnknownValues'] = ignore_unknown_values
+
         try:
             response = table_data.insertAll(
                 projectId=self.project_id,

bigquery/tests/test_client.py

@@ -2104,6 +2104,47 @@ def test_push_success(self):
         self.mock_table_data.insertAll.return_value.execute.assert_has_calls(
             execute_calls)
 
+    def test_request_data_with_options(self):
+        """Ensure that insertAll body has optional property only when
+        the optional parameter of push_rows passed.
+        """
+        expected_body = self.data.copy()
+
+        self.client.push_rows(
+            self.dataset, self.table, self.rows,
+            insert_id_key='one')
+        self.mock_table_data.insertAll.assert_called_with(
+            projectId=self.project,
+            datasetId=self.dataset,
+            tableId=self.table,
+            body=expected_body)
+
+        self.client.push_rows(
+            self.dataset, self.table, self.rows,
+            insert_id_key='one',
+            ignore_unknown_values=False,
+            skip_invalid_rows=False)
+        expected_body['ignoreUnknownValues'] = False
+        expected_body['skipInvalidRows'] = False
+        self.mock_table_data.insertAll.assert_called_with(
+            projectId=self.project,
+            datasetId=self.dataset,
+            tableId=self.table,
+            body=expected_body)
+
+        self.client.push_rows(
+            self.dataset, self.table, self.rows,
+            insert_id_key='one',
+            ignore_unknown_values=True,
+            skip_invalid_rows=True)
+        expected_body['ignoreUnknownValues'] = True
+        expected_body['skipInvalidRows'] = True
+        self.mock_table_data.insertAll.assert_called_with(
+            projectId=self.project,
+            datasetId=self.dataset,
+            tableId=self.table,
+            body=expected_body)
+
 
 class TestGetAllTables(unittest.TestCase):