feat: add LoadJobConfig.projection_fields to select DATASTORE_BACKUP fields (#736)

tswast · plamut · web-flow · commit c45a7380871a · 2021-07-15T11:19:25.000+02:00
* feat: add LoadJobConfig.projection_fields to select DATASTORE_BACKUP fields

* add type annotations

* annotate setter too

Co-authored-by: Peter Lamut &lt;plamut@users.noreply.github.com&gt;
diff --git a/google/cloud/bigquery/job/load.py b/google/cloud/bigquery/job/load.py
@@ -14,7 +14,7 @@
 
 """Classes for load jobs."""
 
-from typing import FrozenSet, Iterable, Optional
+from typing import FrozenSet, List, Iterable, Optional
 
 from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 from google.cloud.bigquery.external_config import HivePartitioningOptions
@@ -25,7 +25,6 @@
 from google.cloud.bigquery.table import RangePartitioning
 from google.cloud.bigquery.table import TableReference
 from google.cloud.bigquery.table import TimePartitioning
-
 from google.cloud.bigquery.job.base import _AsyncJob
 from google.cloud.bigquery.job.base import _JobConfig
 from google.cloud.bigquery.job.base import _JobReference
@@ -300,6 +299,27 @@ def null_marker(self):
     def null_marker(self, value):
         self._set_sub_prop("nullMarker", value)
 
+    @property
+    def projection_fields(self) -> Optional[List[str]]:
+        """Optional[List[str]]: If
+        :attr:`google.cloud.bigquery.job.LoadJobConfig.source_format` is set to
+        "DATASTORE_BACKUP", indicates which entity properties to load into
+        BigQuery from a Cloud Datastore backup.
+
+        Property names are case sensitive and must be top-level properties. If
+        no properties are specified, BigQuery loads all properties. If any
+        named property isn't found in the Cloud Datastore backup, an invalid
+        error is returned in the job result.
+
+        See:
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfigurationLoad.FIELDS.projection_fields
+        """
+        return self._get_sub_prop("projectionFields")
+
+    @projection_fields.setter
+    def projection_fields(self, value: Optional[List[str]]):
+        self._set_sub_prop("projectionFields", value)
+
     @property
     def quote_character(self):
         """Optional[str]: Character used to quote data sections (CSV only).
diff --git a/tests/unit/job/test_load_config.py b/tests/unit/job/test_load_config.py
@@ -424,6 +424,17 @@ def test_null_marker_setter(self):
         config.null_marker = null_marker
         self.assertEqual(config._properties["load"]["nullMarker"], null_marker)
 
+    def test_projection_fields_miss(self):
+        config = self._get_target_class()()
+        self.assertIsNone(config.projection_fields)
+
+    def test_projection_fields_hit(self):
+        config = self._get_target_class()()
+        fields = ["email", "postal_code"]
+        config.projection_fields = fields
+        self.assertEqual(config._properties["load"]["projectionFields"], fields)
+        self.assertEqual(config.projection_fields, fields)
+
     def test_quote_character_missing(self):
         config = self._get_target_class()()
         self.assertIsNone(config.quote_character)
