Merge pull request #128 from cynipe/support-nested-column-for-insert_id

Allow to specify nested column as insertId for push_row

Author: tylertreat

bigquery/client.py

@@ -6,6 +6,7 @@
 from hashlib import sha256
 from io import StringIO
 from time import sleep, time
+from functools import reduce
 
 import six
 from bigquery.errors import (BigQueryTimeoutException, JobExecutingException,
@@ -1236,7 +1237,8 @@ def push_rows(self, dataset, table, rows, insert_id_key=None,
         rows : list
             A ``list`` of rows (``dict`` objects) to add to the table
         insert_id_key : str, optional
-            Key for insertId in row
+            Key for insertId in row.
+            You can use dot separated key for nested column.
         skip_invalid_rows : bool, optional
             Insert all valid rows of a request, even if invalid rows exist.
         ignore_unknown_values : bool, optional
@@ -1258,8 +1260,11 @@ def push_rows(self, dataset, table, rows, insert_id_key=None,
         for row in rows:
             each_row = {}
             each_row["json"] = row
-            if insert_id_key in row:
-                each_row["insertId"] = row[insert_id_key]
+            if insert_id_key is not None:
+                keys = insert_id_key.split('.')
+                val = reduce(lambda d, key: d.get(key) if d else None, keys, row)
+                if val is not None:
+                    each_row["insertId"] = val
             rows_data.append(each_row)
 
         data = {

bigquery/tests/test_client.py

@@ -2325,6 +2325,52 @@ def test_request_data_with_options(self):
             tableId=self.table,
             body=expected_body)
 
+    def test_insert_id_key_with_nested_column(self):
+        """Ensure that dot separated insert_id_key properly extracted with nested column value."""
+        rows = [
+            {'nested': {'col': 'nested_col1'}, 'val': 1},
+            {'nested': {'col': 'nested_col2'}, 'val': 2},
+        ]
+        expected_body = self.data.copy()
+        expected_body['rows'] = [
+            {'insertId': 'nested_col1', 'json': {'nested': {'col': 'nested_col1'}, 'val': 1}},
+            {'insertId': 'nested_col2', 'json': {'nested': {'col': 'nested_col2'}, 'val': 2}},
+        ]
+
+        self.client.push_rows(self.dataset, self.table, rows,
+                              insert_id_key='nested.col')
+        self.mock_table_data.insertAll.assert_called_with(
+            projectId=self.project,
+            datasetId=self.dataset,
+            tableId=self.table,
+            body=expected_body)
+
+        expected_body = self.data.copy()
+        expected_body['rows'] = [
+            {'insertId': 1, 'json': {'nested': {'col': 'nested_col1'}, 'val': 1}},
+            {'insertId': 2, 'json': {'nested': {'col': 'nested_col2'}, 'val': 2}},
+        ]
+        self.client.push_rows(self.dataset, self.table, rows,
+                              insert_id_key='val')
+        self.mock_table_data.insertAll.assert_called_with(
+            projectId=self.project,
+            datasetId=self.dataset,
+            tableId=self.table,
+            body=expected_body)
+
+        expected_body = self.data.copy()
+        expected_body['rows'] = [
+            {'json': {'nested': {'col': 'nested_col1'}, 'val': 1}},
+            {'json': {'nested': {'col': 'nested_col2'}, 'val': 2}},
+        ]
+        self.client.push_rows(self.dataset, self.table, rows,
+                              insert_id_key='no_such.column')
+        self.mock_table_data.insertAll.assert_called_with(
+            projectId=self.project,
+            datasetId=self.dataset,
+            tableId=self.table,
+            body=expected_body)
+
 
 class TestGetAllTables(unittest.TestCase):