feat: support job reservation (#2186)

* feat: support job reservation

* Update google/cloud/bigquery/job/base.py

Co-authored-by: Chalmer Lowe <[email protected]>

---------

Co-authored-by: Chalmer Lowe <[email protected]>

Author: Linchin

google/cloud/bigquery/job/base.py

@@ -224,6 +224,26 @@ def job_timeout_ms(self, value):
         else:
             self._properties.pop("jobTimeoutMs", None)
 
+    @property
+    def reservation(self):
+        """str: Optional. The reservation that job would use.
+
+        User can specify a reservation to execute the job. If reservation is
+        not set, reservation is determined based on the rules defined by the
+        reservation assignments. The expected format is
+        projects/{project}/locations/{location}/reservations/{reservation}.
+
+        Raises:
+            ValueError: If ``value`` type is not None or of string type.
+        """
+        return self._properties.setdefault("reservation", None)
+
+    @reservation.setter
+    def reservation(self, value):
+        if value and not isinstance(value, str):
+            raise ValueError("Reservation must be None or a string.")
+        self._properties["reservation"] = value
+
     @property
     def labels(self):
         """Dict[str, str]: Labels for the job.
@@ -488,6 +508,18 @@ def location(self):
         """str: Location where the job runs."""
         return _helpers._get_sub_prop(self._properties, ["jobReference", "location"])
 
+    @property
+    def reservation_id(self):
+        """str: Name of the primary reservation assigned to this job.
+
+        Note that this could be different than reservations reported in
+        the reservation field if parent reservations were used to execute
+        this job.
+        """
+        return _helpers._get_sub_prop(
+            self._properties, ["statistics", "reservation_id"]
+        )
+
     def _require_client(self, client):
         """Check client or verify over-ride.
 

tests/unit/job/test_base.py

@@ -443,6 +443,16 @@ def test_state(self):
         status["state"] = state
         self.assertEqual(job.state, state)
 
+    def test_reservation_id(self):
+        reservation_id = "RESERVATION-ID"
+        client = _make_client(project=self.PROJECT)
+        job = self._make_one(self.JOB_ID, client)
+        self.assertIsNone(job.reservation_id)
+        stats = job._properties["statistics"] = {}
+        self.assertIsNone(job.reservation_id)
+        stats["reservation_id"] = reservation_id
+        self.assertEqual(job.reservation_id, reservation_id)
+
     def _set_properties_job(self):
         client = _make_client(project=self.PROJECT)
         job = self._make_one(self.JOB_ID, client)
@@ -1188,31 +1198,37 @@ def test_fill_query_job_config_from_default(self):
         job_config = QueryJobConfig()
         job_config.dry_run = True
         job_config.maximum_bytes_billed = 1000
+        job_config.reservation = "reservation_1"
 
         default_job_config = QueryJobConfig()
         default_job_config.use_query_cache = True
         default_job_config.maximum_bytes_billed = 2000
+        default_job_config.reservation = "reservation_2"
 
         final_job_config = job_config._fill_from_default(default_job_config)
         self.assertTrue(final_job_config.dry_run)
         self.assertTrue(final_job_config.use_query_cache)
         self.assertEqual(final_job_config.maximum_bytes_billed, 1000)
+        self.assertEqual(final_job_config.reservation, "reservation_1")
 
     def test_fill_load_job_from_default(self):
         from google.cloud.bigquery import LoadJobConfig
 
         job_config = LoadJobConfig()
         job_config.create_session = True
         job_config.encoding = "UTF-8"
+        job_config.reservation = "reservation_1"
 
         default_job_config = LoadJobConfig()
         default_job_config.ignore_unknown_values = True
         default_job_config.encoding = "ISO-8859-1"
+        default_job_config.reservation = "reservation_2"
 
         final_job_config = job_config._fill_from_default(default_job_config)
         self.assertTrue(final_job_config.create_session)
         self.assertTrue(final_job_config.ignore_unknown_values)
         self.assertEqual(final_job_config.encoding, "UTF-8")
+        self.assertEqual(final_job_config.reservation, "reservation_1")
 
     def test_fill_from_default_conflict(self):
         from google.cloud.bigquery import QueryJobConfig
@@ -1232,10 +1248,12 @@ def test_fill_from_empty_default_conflict(self):
         job_config = QueryJobConfig()
         job_config.dry_run = True
         job_config.maximum_bytes_billed = 1000
+        job_config.reservation = "reservation_1"
 
         final_job_config = job_config._fill_from_default(default_job_config=None)
         self.assertTrue(final_job_config.dry_run)
         self.assertEqual(final_job_config.maximum_bytes_billed, 1000)
+        self.assertEqual(final_job_config.reservation, "reservation_1")
 
     @mock.patch("google.cloud.bigquery._helpers._get_sub_prop")
     def test__get_sub_prop_wo_default(self, _get_sub_prop):
@@ -1338,3 +1356,27 @@ def test_job_timeout_properties(self):
         job_config.job_timeout_ms = None
         assert job_config.job_timeout_ms is None
         assert "jobTimeoutMs" not in job_config._properties
+
+    def test_reservation_miss(self):
+        job_config = self._make_one()
+        self.assertEqual(job_config.reservation, None)
+
+    def test_reservation_hit(self):
+        job_config = self._make_one()
+        job_config._properties["reservation"] = "foo"
+        self.assertEqual(job_config.reservation, "foo")
+
+    def test_reservation_update_in_place(self):
+        job_config = self._make_one()
+        job_config.reservation = "bar"  # update in place
+        self.assertEqual(job_config.reservation, "bar")
+
+    def test_reservation_setter_invalid(self):
+        job_config = self._make_one()
+        with self.assertRaises(ValueError):
+            job_config.reservation = object()
+
+    def test_reservation_setter(self):
+        job_config = self._make_one()
+        job_config.reservation = "foo"
+        self.assertEqual(job_config._properties["reservation"], "foo")