Add avro logical type control for load jobs. (googleapis#6827)

shollyman · tseaver · commit d664d2dd392e · 2018-12-03T14:53:06.000-05:00
diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
@@ -1195,6 +1195,18 @@ def time_partitioning(self, value):
         else:
             self._del_sub_prop("timePartitioning")
 
+    @property
+    def use_avro_logical_types(self):
+        """bool: For loads of Avro data, governs whether Avro logical types are
+        converted to their corresponding BigQuery types(e.g. TIMESTAMP) rather than
+        raw types (e.g. INTEGER).
+        """
+        return self._get_sub_prop("useAvroLogicalTypes")
+
+    @use_avro_logical_types.setter
+    def use_avro_logical_types(self, value):
+        self._set_sub_prop("useAvroLogicalTypes", bool(value))
+
     @property
     def write_disposition(self):
         """google.cloud.bigquery.job.WriteDisposition: Action that occurs if
@@ -1362,6 +1374,13 @@ def time_partitioning(self):
         """
         return self._configuration.time_partitioning
 
+    @property
+    def use_avro_logical_types(self):
+        """See
+        :attr:`google.cloud.bigquery.job.LoadJobConfig.use_avro_logical_types`.
+        """
+        return self._configuration.use_avro_logical_types
+
     @property
     def clustering_fields(self):
         """See
diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
@@ -1588,6 +1588,15 @@ def test_time_partitioning_setter_w_none(self):
         self.assertIsNone(config.time_partitioning)
         self.assertNotIn("timePartitioning", config._properties["load"])
 
+    def test_use_avro_logical_types(self):
+        config = self._get_target_class()()
+        self.assertIsNone(config.use_avro_logical_types)
+
+    def test_use_avro_logical_types_setter(self):
+        config = self._get_target_class()()
+        config.use_avro_logical_types = True
+        self.assertTrue(config._properties["load"]["useAvroLogicalTypes"])
+
     def test_write_disposition_missing(self):
         config = self._get_target_class()()
         self.assertIsNone(config.write_disposition)
@@ -1663,6 +1672,10 @@ def _verifyBooleanConfigProperties(self, job, config):
             self.assertEqual(job.ignore_unknown_values, config["ignoreUnknownValues"])
         else:
             self.assertIsNone(job.ignore_unknown_values)
+        if "useAvroLogicalTypes" in config:
+            self.assertEqual(job.use_avro_logical_types, config["useAvroLogicalTypes"])
+        else:
+            self.assertIsNone(job.use_avro_logical_types)
 
     def _verifyEnumConfigProperties(self, job, config):
         if "createDisposition" in config:
@@ -1765,6 +1778,7 @@ def test_ctor(self):
         self.assertIsNone(job.write_disposition)
         self.assertIsNone(job.destination_encryption_configuration)
         self.assertIsNone(job.time_partitioning)
+        self.assertIsNone(job.use_avro_logical_types)
         self.assertIsNone(job.clustering_fields)
         self.assertIsNone(job.schema_update_options)
 
@@ -2097,6 +2111,7 @@ def test_begin_w_alternate_client(self):
             "quote": "'",
             "skipLeadingRows": "1",
             "sourceFormat": "CSV",
+            "useAvroLogicalTypes": True,
             "writeDisposition": WriteDisposition.WRITE_TRUNCATE,
             "schema": {
                 "fields": [
@@ -2139,6 +2154,7 @@ def test_begin_w_alternate_client(self):
         config.quote_character = "'"
         config.skip_leading_rows = 1
         config.source_format = "CSV"
+        config.use_avro_logical_types = True
         config.write_disposition = WriteDisposition.WRITE_TRUNCATE
         config.schema_update_options = [SchemaUpdateOption.ALLOW_FIELD_ADDITION]
 
