feat: Add python client api for JobStatistics.referenced_property_graphs.

Author: ericfe-google

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.

packages/google-cloud-bigquery/google/cloud/bigquery/table.py

@@ -138,7 +138,7 @@ def _reference_getter(table):
 
 
 def _view_use_legacy_sql_getter(
-    table: Union["Table", "TableListItem"]
+    table: Union["Table", "TableListItem"],
 ) -> Optional[bool]:
     """bool: Specifies whether to execute the view with Legacy or Standard SQL.
 
@@ -359,6 +359,131 @@ def __repr__(self):
         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.
 
@@ -452,9 +577,9 @@ def biglake_configuration(self, value):
         api_repr = value
         if value is not None:
             api_repr = value.to_api_repr()
-        self._properties[
-            self._PROPERTY_TO_API_FIELD["biglake_configuration"]
-        ] = api_repr
+        self._properties[self._PROPERTY_TO_API_FIELD["biglake_configuration"]] = (
+            api_repr
+        )
 
     @property
     def require_partition_filter(self):
@@ -468,9 +593,9 @@ def require_partition_filter(self):
 
     @require_partition_filter.setter
     def require_partition_filter(self, value):
-        self._properties[
-            self._PROPERTY_TO_API_FIELD["require_partition_filter"]
-        ] = value
+        self._properties[self._PROPERTY_TO_API_FIELD["require_partition_filter"]] = (
+            value
+        )
 
     @property
     def schema(self):
@@ -568,9 +693,9 @@ def encryption_configuration(self, value):
         api_repr = value
         if value is not None:
             api_repr = value.to_api_repr()
-        self._properties[
-            self._PROPERTY_TO_API_FIELD["encryption_configuration"]
-        ] = api_repr
+        self._properties[self._PROPERTY_TO_API_FIELD["encryption_configuration"]] = (
+            api_repr
+        )
 
     @property
     def created(self):
@@ -845,9 +970,9 @@ def expires(self, value):
         if not isinstance(value, datetime.datetime) and value is not None:
             raise ValueError("Pass a datetime, or None")
         value_ms = google.cloud._helpers._millis_from_datetime(value)
-        self._properties[
-            self._PROPERTY_TO_API_FIELD["expires"]
-        ] = _helpers._str_or_none(value_ms)
+        self._properties[self._PROPERTY_TO_API_FIELD["expires"]] = (
+            _helpers._str_or_none(value_ms)
+        )
 
     @property
     def friendly_name(self):
@@ -1043,9 +1168,9 @@ def external_data_configuration(self, value):
         api_repr = value
         if value is not None:
             api_repr = value.to_api_repr()
-        self._properties[
-            self._PROPERTY_TO_API_FIELD["external_data_configuration"]
-        ] = api_repr
+        self._properties[self._PROPERTY_TO_API_FIELD["external_data_configuration"]] = (
+            api_repr
+        )
 
     @property
     def snapshot_definition(self) -> Optional["SnapshotDefinition"]:

packages/google-cloud-bigquery/tests/unit/job/test_query.py

@@ -680,6 +680,55 @@ 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 = [
             {
@@ -1586,12 +1635,10 @@ def test_result_with_start_index_multi_page(self):
     def test_result_error(self):
         from google.cloud import exceptions
 
-        query = textwrap.dedent(
-            """
+        query = textwrap.dedent("""
             SELECT foo, bar
             FROM table_baz
-            WHERE foo == bar"""
-        )
+            WHERE foo == bar""")
 
         client = _make_client(project=self.PROJECT)
         job = self._make_one(self.JOB_ID, query, client)
@@ -1635,12 +1682,10 @@ def test_result_error(self):
             assert expected_line in debug_message
 
     def test_result_transport_timeout_error(self):
-        query = textwrap.dedent(
-            """
+        query = textwrap.dedent("""
             SELECT foo, bar
             FROM table_baz
-            WHERE foo == bar"""
-        )
+            WHERE foo == bar""")
 
         client = _make_client(project=self.PROJECT)
         job = self._make_one(self.JOB_ID, query, client)
@@ -1694,12 +1739,10 @@ def test_schema(self):
     def test__begin_error(self):
         from google.cloud import exceptions
 
-        query = textwrap.dedent(
-            """
+        query = textwrap.dedent("""
             SELECT foo, bar
             FROM table_baz
-            WHERE foo == bar"""
-        )
+            WHERE foo == bar""")
 
         client = _make_client(project=self.PROJECT)
         job = self._make_one(self.JOB_ID, query, client)