From d62cabbf115637ecbaf8cc378f39329a5ae74c26 Mon Sep 17 00:00:00 2001
From: Lingqing Gan <lingqing.gan@gmail.com>
Date: Tue, 27 Feb 2024 12:18:10 -0800
Subject: [PATCH] feat: support slot_ms in QueryPlanEntry (#1831)

---
 google/cloud/bigquery/job/query.py | 5 +++++
 tests/unit/job/test_query_stats.py | 4 ++++
 2 files changed, 9 insertions(+)

diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py
index ac0c51973..e45a46894 100644
--- a/google/cloud/bigquery/job/query.py
+++ b/google/cloud/bigquery/job/query.py
@@ -2263,6 +2263,11 @@ def steps(self):
             for step in self._properties.get("steps", [])
         ]
 
+    @property
+    def slot_ms(self):
+        """Optional[int]: Slot-milliseconds used by the stage."""
+        return _helpers._int_or_none(self._properties.get("slotMs"))
+
 
 class TimelineEntry(object):
     """TimelineEntry represents progress of a query job at a particular
diff --git a/tests/unit/job/test_query_stats.py b/tests/unit/job/test_query_stats.py
index bdd0fb627..61b278d43 100644
--- a/tests/unit/job/test_query_stats.py
+++ b/tests/unit/job/test_query_stats.py
@@ -261,6 +261,7 @@ class TestQueryPlanEntry(_Base):
     STATUS = "STATUS"
     SHUFFLE_OUTPUT_BYTES = 1024
     SHUFFLE_OUTPUT_BYTES_SPILLED = 1
+    SLOT_MS = 25
 
     START_RFC3339_MICROS = "2018-04-01T00:00:00.000000Z"
     END_RFC3339_MICROS = "2018-04-01T00:00:04.000000Z"
@@ -305,6 +306,7 @@ def test_from_api_repr_empty(self):
         self.assertIsNone(entry.shuffle_output_bytes)
         self.assertIsNone(entry.shuffle_output_bytes_spilled)
         self.assertEqual(entry.steps, [])
+        self.assertIsNone(entry.slot_ms)
 
     def test_from_api_repr_normal(self):
         from google.cloud.bigquery.job import QueryPlanEntryStep
@@ -348,6 +350,7 @@ def test_from_api_repr_normal(self):
                     "substeps": TestQueryPlanEntryStep.SUBSTEPS,
                 }
             ],
+            "slotMs": self.SLOT_MS,
         }
         klass = self._get_target_class()
 
@@ -366,6 +369,7 @@ def test_from_api_repr_normal(self):
         self.assertEqual(entry.records_written, self.RECORDS_WRITTEN)
         self.assertEqual(entry.status, self.STATUS)
         self.assertEqual(entry.steps, steps)
+        self.assertEqual(entry.slot_ms, self.SLOT_MS)
 
     def test_start(self):
         from google.cloud._helpers import _RFC3339_MICROS