feat: Add python client api for JobStatistics.referenced_property_graphs.

ericfe-google · ericfe-google · commit f4aecbc2bb14 · 2026-03-20T16:38:35.000-07:00
diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py
@@ -50,7 +50,7 @@
 from google.cloud.bigquery.table import _EmptyRowIterator
 from google.cloud.bigquery.table import RangePartitioning
 from google.cloud.bigquery.table import _table_arg_to_table_ref
-from google.cloud.bigquery.table import TableReference
+from google.cloud.bigquery.table import TableReference, PropertyGraphReference
 from google.cloud.bigquery.table import TimePartitioning
 from google.cloud.bigquery._tqdm_helpers import wait_for_query
 
@@ -1332,6 +1332,35 @@ def referenced_tables(self):
 
         return tables
 
+    @property
+    def referenced_property_graphs(self):
+        """Return referenced property graphs from job statistics, if present.
+
+        See:
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobStatistics2.FIELDS.referenced_property_graphs
+
+        Returns:
+            List[google.cloud.bigquery.table.PropertyGraphReference]:
+                mappings describing the property graphs, or an empty list
+                if the query has not yet completed.
+        """
+        property_graphs = []
+        datasets_by_project_name = {}
+
+        for pg in self._job_statistics().get("referencedPropertyGraphs", ()):
+            pg_project = pg["projectId"]
+
+            ds_id = pg["datasetId"]
+            pg_dataset = datasets_by_project_name.get((pg_project, ds_id))
+            if pg_dataset is None:
+                pg_dataset = DatasetReference(pg_project, ds_id)
+                datasets_by_project_name[(pg_project, ds_id)] = pg_dataset
+
+            pg_name = pg["propertyGraphId"]
+            property_graphs.append(PropertyGraphReference(pg_dataset, pg_name))
+
+        return property_graphs
+
     @property
     def undeclared_query_parameters(self):
         """Return undeclared query parameters from job statistics, if present.
diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py
@@ -358,6 +358,127 @@ def __repr__(self):
         dataset_ref = DatasetReference(self.project, self.dataset_id)
         return f"TableReference({dataset_ref!r}, '{self.table_id}')"
 
+class PropertyGraphReference:
+    """PropertyGraphReferences are pointers to property graphs.
+
+    Args:
+        dataset_ref: A pointer to the dataset
+        property_graph_id: The ID of the property graph
+    """
+
+    _PROPERTY_TO_API_FIELD = {
+        "dataset_id": "datasetId",
+        "project": "projectId",
+        "property_graph_id": "propertyGraphId",
+    }
+
+    def __init__(self, dataset_ref: "DatasetReference", property_graph_id: str):
+        self._properties = {}
+
+        _helpers._set_sub_prop(
+            self._properties,
+            self._PROPERTY_TO_API_FIELD["project"],
+            dataset_ref.project,
+        )
+        _helpers._set_sub_prop(
+            self._properties,
+            self._PROPERTY_TO_API_FIELD["dataset_id"],
+            dataset_ref.dataset_id,
+        )
+        _helpers._set_sub_prop(
+            self._properties,
+            self._PROPERTY_TO_API_FIELD["property_graph_id"],
+            property_graph_id,
+        )
+
+    @property
+    def project(self) -> str:
+        """str: Project bound to the property graph."""
+        return _helpers._get_sub_prop(
+            self._properties, self._PROPERTY_TO_API_FIELD["project"]
+        )
+
+    @property
+    def dataset_id(self) -> str:
+        """str: ID of dataset containing the property graph."""
+        return _helpers._get_sub_prop(
+            self._properties, self._PROPERTY_TO_API_FIELD["dataset_id"]
+        )
+
+    @property
+    def property_graph_id(self) -> str:
+        """str: The property graph ID."""
+        return _helpers._get_sub_prop(
+            self._properties, self._PROPERTY_TO_API_FIELD["property_graph_id"]
+        )
+
+    @classmethod
+    def from_string(
+        cls, property_graph_id: str, default_project: Optional[str] = None
+    ) -> "PropertyGraphReference":
+        """Construct a property graph reference from string.
+
+        Args:
+            property_graph_id (str):
+                A property graph ID in standard SQL format.
+            default_project (Optional[str]):
+                The project ID to use when ``property_graph_id`` does not
+                include a project ID.
+
+        Returns:
+            PropertyGraphReference: Property graph reference parsed from ``property_graph_id``.
+        """
+        from google.cloud.bigquery.dataset import DatasetReference
+
+        (
+            output_project_id,
+            output_dataset_id,
+            output_property_graph_id,
+        ) = _helpers._parse_3_part_id(
+            property_graph_id, default_project=default_project, property_name="property_graph_id"
+        )
+
+        return cls(
+            DatasetReference(output_project_id, output_dataset_id), output_property_graph_id
+        )
+
+    @classmethod
+    def from_api_repr(cls, resource: dict) -> "PropertyGraphReference":
+        """Factory: construct a property graph reference given its API representation."""
+        from google.cloud.bigquery.dataset import DatasetReference
+
+        project = resource["projectId"]
+        dataset_id = resource["datasetId"]
+        property_graph_id = resource["propertyGraphId"]
+
+        return cls(DatasetReference(project, dataset_id), property_graph_id)
+
+    def to_api_repr(self) -> dict:
+        """Construct the API resource representation of this property graph reference."""
+        return copy.deepcopy(self._properties)
+
+    def __str__(self):
+        return f"{self.project}.{self.dataset_id}.{self.property_graph_id}"
+
+    def __repr__(self):
+        from google.cloud.bigquery.dataset import DatasetReference
+
+        dataset_ref = DatasetReference(self.project, self.dataset_id)
+        return f"PropertyGraphReference({dataset_ref!r}, '{self.property_graph_id}')"
+
+    def __eq__(self, other):
+        if isinstance(other, PropertyGraphReference):
+            return (
+                self.project == other.project
+                and self.dataset_id == other.dataset_id
+                and self.property_graph_id == other.property_graph_id
+            )
+        else:
+            return NotImplemented
+
+    def __hash__(self):
+        return hash((self.project, self.dataset_id, self.property_graph_id))
+
 
 class Table(_TableBase):
     """Tables represent a set of rows whose values correspond to a schema.
diff --git a/packages/google-cloud-bigquery/tests/unit/job/test_query.py b/packages/google-cloud-bigquery/tests/unit/job/test_query.py
@@ -680,6 +680,47 @@ def test_referenced_tables(self):
         self.assertEqual(remote.dataset_id, "other-dataset")
         self.assertEqual(remote.project, "other-project-123")
 
+    def test_referenced_property_graphs(self):
+        from google.cloud.bigquery.table import PropertyGraphReference
+
+        ref_pg_resource = [
+            {"projectId": self.PROJECT, "datasetId": "dataset", "propertyGraphId": "pg1"},
+            {"projectId": self.PROJECT, "datasetId": "dataset", "propertyGraphId": "pg2"},
+            {
+                "projectId": "other-project-123",
+                "datasetId": "other-dataset",
+                "propertyGraphId": "other-pg",
+            },
+        ]
+        client = _make_client(project=self.PROJECT)
+        job = self._make_one(self.JOB_ID, self.QUERY, client)
+        self.assertEqual(job.referenced_property_graphs, [])
+
+        statistics = job._properties["statistics"] = {}
+        self.assertEqual(job.referenced_property_graphs, [])
+
+        query_stats = statistics["query"] = {}
+        self.assertEqual(job.referenced_property_graphs, [])
+
+        query_stats["referencedPropertyGraphs"] = ref_pg_resource
+
+        pg1, pg2, remote = job.referenced_property_graphs
+
+        self.assertIsInstance(pg1, PropertyGraphReference)
+        self.assertEqual(pg1.property_graph_id, "pg1")
+        self.assertEqual(pg1.dataset_id, "dataset")
+        self.assertEqual(pg1.project, self.PROJECT)
+
+        self.assertIsInstance(pg2, PropertyGraphReference)
+        self.assertEqual(pg2.property_graph_id, "pg2")
+        self.assertEqual(pg2.dataset_id, "dataset")
+        self.assertEqual(pg2.project, self.PROJECT)
+
+        self.assertIsInstance(remote, PropertyGraphReference)
+        self.assertEqual(remote.property_graph_id, "other-pg")
+        self.assertEqual(remote.dataset_id, "other-dataset")
+        self.assertEqual(remote.project, "other-project-123")
+
     def test_timeline(self):
         timeline_resource = [
             {
