Add 'Table.row_from_mapping' helper. (googleapis#3425)

tseaver · lukesneeringer · commit 1b764b80e2aa · 2017-08-07T14:50:10.000-07:00
diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py
@@ -733,6 +733,35 @@ def fetch_data(self, max_results=None, page_token=None, client=None):
         iterator._NEXT_TOKEN = 'pageToken'
         return iterator
 
+    def row_from_mapping(self, mapping):
+        """Convert a mapping to a row tuple using the schema.
+
+        :type mapping: dict
+        :param mapping: Mapping of row data: must contain keys for all
+               required fields in the schema.  Keys which do not correspond
+               to a field in the schema are ignored.
+
+        :rtype: tuple
+        :returns: Tuple whose elements are ordered according to the table's
+                  schema.
+        :raises: ValueError if table's schema is not set
+        """
+        if len(self._schema) == 0:
+            raise ValueError(_TABLE_HAS_NO_SCHEMA)
+
+        row = []
+        for field in self.schema:
+            if field.mode == 'REQUIRED':
+                row.append(mapping[field.name])
+            elif field.mode == 'REPEATED':
+                row.append(mapping.get(field.name, ()))
+            elif field.mode == 'NULLABLE':
+                row.append(mapping.get(field.name))
+            else:
+                raise ValueError(
+                    "Unknown field mode: {}".format(field.mode))
+        return tuple(row)
+
     def insert_data(self,
                     rows,
                     row_ids=None,
diff --git a/bigquery/tests/unit/test_table.py b/bigquery/tests/unit/test_table.py
@@ -1347,6 +1347,61 @@ def test_fetch_data_w_record_schema(self):
         self.assertEqual(req['method'], 'GET')
         self.assertEqual(req['path'], '/%s' % PATH)
 
+    def test_row_from_mapping_wo_schema(self):
+        from google.cloud.bigquery.table import _TABLE_HAS_NO_SCHEMA
+        MAPPING = {'full_name': 'Phred Phlyntstone', 'age': 32}
+        client = _Client(project=self.PROJECT)
+        dataset = _Dataset(client)
+        table = self._make_one(self.TABLE_NAME, dataset=dataset)
+
+        with self.assertRaises(ValueError) as exc:
+            table.row_from_mapping(MAPPING)
+
+        self.assertEqual(exc.exception.args, (_TABLE_HAS_NO_SCHEMA,))
+
+    def test_row_from_mapping_w_invalid_schema(self):
+        from google.cloud.bigquery.table import SchemaField
+        MAPPING = {
+            'full_name': 'Phred Phlyntstone',
+            'age': 32,
+            'colors': ['red', 'green'],
+            'bogus': 'WHATEVER',
+        }
+        client = _Client(project=self.PROJECT)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        colors = SchemaField('colors', 'DATETIME', mode='REPEATED')
+        bogus = SchemaField('joined', 'STRING', mode='BOGUS')
+        table = self._make_one(self.TABLE_NAME, dataset=dataset,
+                               schema=[full_name, age, colors, bogus])
+
+        with self.assertRaises(ValueError) as exc:
+            table.row_from_mapping(MAPPING)
+
+        self.assertIn('Unknown field mode: BOGUS', str(exc.exception))
+
+    def test_row_from_mapping_w_schema(self):
+        from google.cloud.bigquery.table import SchemaField
+        MAPPING = {
+            'full_name': 'Phred Phlyntstone',
+            'age': 32,
+            'colors': ['red', 'green'],
+            'extra': 'IGNORED',
+        }
+        client = _Client(project=self.PROJECT)
+        dataset = _Dataset(client)
+        full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
+        age = SchemaField('age', 'INTEGER', mode='REQUIRED')
+        colors = SchemaField('colors', 'DATETIME', mode='REPEATED')
+        joined = SchemaField('joined', 'STRING', mode='NULLABLE')
+        table = self._make_one(self.TABLE_NAME, dataset=dataset,
+                               schema=[full_name, age, colors, joined])
+
+        self.assertEqual(
+            table.row_from_mapping(MAPPING),
+            ('Phred Phlyntstone', 32, ['red', 'green'], None))
+
     def test_insert_data_wo_schema(self):
         from google.cloud.bigquery.table import _TABLE_HAS_NO_SCHEMA
 
@@ -2055,7 +2110,7 @@ def test__parse_schema_resource_subfields(self):
         RESOURCE['schema']['fields'].append(
             {'name': 'phone',
              'type': 'RECORD',
-             'mode': 'REPEATABLE',
+             'mode': 'REPEATED',
              'fields': [{'name': 'type',
                          'type': 'STRING',
                          'mode': 'REQUIRED'},
@@ -2123,7 +2178,7 @@ def test_w_subfields(self):
         full_name = SchemaField('full_name', 'STRING', mode='REQUIRED')
         ph_type = SchemaField('type', 'STRING', 'REQUIRED')
         ph_num = SchemaField('number', 'STRING', 'REQUIRED')
-        phone = SchemaField('phone', 'RECORD', mode='REPEATABLE',
+        phone = SchemaField('phone', 'RECORD', mode='REPEATED',
                             fields=[ph_type, ph_num])
         resource = self._call_fut([full_name, phone])
         self.assertEqual(len(resource), 2)
@@ -2134,7 +2189,7 @@ def test_w_subfields(self):
         self.assertEqual(resource[1],
                          {'name': 'phone',
                           'type': 'RECORD',
-                          'mode': 'REPEATABLE',
+                          'mode': 'REPEATED',
                           'fields': [{'name': 'type',
                                       'type': 'STRING',
                                       'mode': 'REQUIRED'},
