pass timestamp_parser to schema_from_record

scribu · scribu · commit 93c1c3570c21 · 2015-07-01T17:43:00.000+03:00
diff --git a/bigquery/schema_builder.py b/bigquery/schema_builder.py
@@ -75,7 +75,7 @@ def bq_schema_field(name, bq_type, mode):
     field = bq_schema_field(k, bq_type, mode)
     if bq_type == "record":
         try:
-            field['fields'] = schema_from_record(v)
+            field['fields'] = schema_from_record(v, timestamp_parser)
         except InvalidTypeException, e:
             # recursively construct the key causing the error
             raise InvalidTypeException("%s.%s" % (k, e.key), e.value)
diff --git a/bigquery/tests/test_schema_builder.py b/bigquery/tests/test_schema_builder.py
@@ -90,6 +90,33 @@ def test_hierarchical_record(self):
 
         self.assertItemsEqual(schema_from_record(record), schema)
 
+    def test_hierarchical_record_with_timestamps(self):
+        record = {"global": "2001-01-01", "user": {"local": "2001-01-01"}}
+
+        schema_with_ts = [
+            {"name": "global", "type": "timestamp", "mode": "nullable"},
+            {"name": "user", "type": "record", "mode": "nullable",
+                "fields": [{
+                    "name": "local",
+                    "type": "timestamp",
+                    "mode": "nullable"}]}]
+
+        schema_without_ts = [
+            {"name": "global", "type": "string", "mode": "nullable"},
+            {"name": "user", "type": "record", "mode": "nullable",
+                "fields": [{
+                    "name": "local",
+                    "type": "string",
+                    "mode": "nullable"}]}]
+
+        self.assertItemsEqual(
+            schema_from_record(record),
+            schema_with_ts)
+
+        self.assertItemsEqual(
+            schema_from_record(record, timestamp_parser=lambda x: False),
+            schema_without_ts)
+
     def test_repeated_field(self):
         record = {"ids": [1, 2, 3, 4, 5]}
         schema = [{"name": "ids", "type": "integer", "mode": "repeated"}]
