Skip to content

Commit 1b4467f

Browse files
ericfe-googleericfe-google
authored
feat: Add python client api for JobStatistics.referenced_property_graphs. (#16137)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/google-cloud-python/issues) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent 1eb7c26 commit 1b4467f

File tree

5 files changed

+172
-7
lines changed

5 files changed

+172
-7
lines changed

‎packages/google-cloud-bigquery/google/cloud/bigquery/job/query.py‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@
5050
from google.cloud.bigquery.table import _EmptyRowIterator
5151
from google.cloud.bigquery.table import RangePartitioning
5252
from google.cloud.bigquery.table import _table_arg_to_table_ref
53-
from google.cloud.bigquery.table import TableReference
53+
from google.cloud.bigquery.table import TableReference, PropertyGraphReference
5454
from google.cloud.bigquery.table import TimePartitioning
5555
from google.cloud.bigquery._tqdm_helpers import wait_for_query
5656

@@ -1332,6 +1332,30 @@ def referenced_tables(self):
13321332

13331333
return tables
13341334

1335+
@property
1336+
def referenced_property_graphs(self):
1337+
"""Return referenced property graphs from job statistics, if present.
1338+
1339+
See:
1340+
https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobStatistics2.FIELDS.referenced_property_graphs
1341+
1342+
Returns:
1343+
List[google.cloud.bigquery.table.PropertyGraphReference]:
1344+
mappings describing the property graphs, or an empty list
1345+
if the query has not yet completed.
1346+
"""
1347+
property_graphs = []
1348+
1349+
for pg in self._job_statistics().get("referencedPropertyGraphs", ()):
1350+
property_graphs.append(
1351+
PropertyGraphReference(
1352+
DatasetReference(pg["projectId"], pg["datasetId"]),
1353+
pg["propertyGraphId"],
1354+
)
1355+
)
1356+
1357+
return property_graphs
1358+
13351359
@property
13361360
def undeclared_query_parameters(self):
13371361
"""Return undeclared query parameters from job statistics, if present.

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

Lines changed: 93 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ def _reference_getter(table):
138138

139139

140140
def _view_use_legacy_sql_getter(
141-
table: Union["Table", "TableListItem"]
141+
table: Union["Table", "TableListItem"],
142142
) -> Optional[bool]:
143143
"""bool: Specifies whether to execute the view with Legacy or Standard SQL.
144144
@@ -359,6 +359,98 @@ def __repr__(self):
359359
return f"TableReference({dataset_ref!r}, '{self.table_id}')"
360360

361361

362+
class PropertyGraphReference:
363+
"""PropertyGraphReferences are pointers to property graphs.
364+
365+
Args:
366+
dataset_ref: A pointer to the dataset
367+
property_graph_id: The ID of the property graph
368+
"""
369+
370+
_PROPERTY_TO_API_FIELD = {
371+
"dataset_id": "datasetId",
372+
"project": "projectId",
373+
"property_graph_id": "propertyGraphId",
374+
}
375+
376+
def __init__(self, dataset_ref: "DatasetReference", property_graph_id: str):
377+
self._properties: Dict[str, Any] = {}
378+
379+
_helpers._set_sub_prop(
380+
self._properties,
381+
self._PROPERTY_TO_API_FIELD["project"],
382+
dataset_ref.project,
383+
)
384+
_helpers._set_sub_prop(
385+
self._properties,
386+
self._PROPERTY_TO_API_FIELD["dataset_id"],
387+
dataset_ref.dataset_id,
388+
)
389+
_helpers._set_sub_prop(
390+
self._properties,
391+
self._PROPERTY_TO_API_FIELD["property_graph_id"],
392+
property_graph_id,
393+
)
394+
395+
@property
396+
def project(self) -> str:
397+
"""str: Project bound to the property graph."""
398+
return _helpers._get_sub_prop(
399+
self._properties, self._PROPERTY_TO_API_FIELD["project"]
400+
)
401+
402+
@property
403+
def dataset_id(self) -> str:
404+
"""str: ID of dataset containing the property graph."""
405+
return _helpers._get_sub_prop(
406+
self._properties, self._PROPERTY_TO_API_FIELD["dataset_id"]
407+
)
408+
409+
@property
410+
def property_graph_id(self) -> str:
411+
"""str: The property graph ID."""
412+
return _helpers._get_sub_prop(
413+
self._properties, self._PROPERTY_TO_API_FIELD["property_graph_id"]
414+
)
415+
416+
@classmethod
417+
def from_api_repr(cls, resource: dict) -> "PropertyGraphReference":
418+
"""Factory: construct a property graph reference given its API representation."""
419+
from google.cloud.bigquery.dataset import DatasetReference
420+
421+
project = resource["projectId"]
422+
dataset_id = resource["datasetId"]
423+
property_graph_id = resource["propertyGraphId"]
424+
425+
return cls(DatasetReference(project, dataset_id), property_graph_id)
426+
427+
def to_api_repr(self) -> dict:
428+
"""Construct the API resource representation of this property graph reference."""
429+
return copy.deepcopy(self._properties)
430+
431+
def __str__(self):
432+
return f"{self.project}.{self.dataset_id}.{self.property_graph_id}"
433+
434+
def __repr__(self):
435+
from google.cloud.bigquery.dataset import DatasetReference
436+
437+
dataset_ref = DatasetReference(self.project, self.dataset_id)
438+
return f"PropertyGraphReference({dataset_ref!r}, '{self.property_graph_id}')"
439+
440+
def __eq__(self, other):
441+
if isinstance(other, PropertyGraphReference):
442+
return (
443+
self.project == other.project
444+
and self.dataset_id == other.dataset_id
445+
and self.property_graph_id == other.property_graph_id
446+
)
447+
else:
448+
return NotImplemented
449+
450+
def __hash__(self):
451+
return hash((self.project, self.dataset_id, self.property_graph_id))
452+
453+
362454
class Table(_TableBase):
363455
"""Tables represent a set of rows whose values correspond to a schema.
364456

‎packages/google-cloud-bigquery/noxfile.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,6 @@ def wrapper(*args, **kwargs):
6868
# 'docfx' is excluded since it only needs to run in 'docs-presubmit'
6969
nox.options.sessions = [
7070
"unit",
71-
"unit_noextras",
7271
"mypy",
7372
"system",
7473
"snippets",

‎packages/google-cloud-bigquery/tests/unit/job/test_async_job_retry.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,6 @@
2626

2727

2828
def test_cancel_w_custom_retry(global_time_lock):
29-
from google.cloud.bigquery.retry import DEFAULT_RETRY
30-
3129
api_path = "/projects/{}/jobs/{}/cancel".format(PROJECT, JOB_ID)
3230
resource = {
3331
"jobReference": {
@@ -49,8 +47,11 @@ def test_cancel_w_custom_retry(global_time_lock):
4947
google.cloud.bigquery.job._JobReference(JOB_ID, PROJECT, "EU"), client
5048
)
5149

52-
retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
53-
lambda exc: isinstance(exc, ValueError)
50+
retry = google.api_core.retry.Retry(
51+
predicate=lambda exc: isinstance(exc, ValueError),
52+
initial=0.01,
53+
maximum=0.01,
54+
deadline=1.0,
5455
)
5556

5657
with mock.patch(

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -680,6 +680,55 @@ def test_referenced_tables(self):
680680
self.assertEqual(remote.dataset_id, "other-dataset")
681681
self.assertEqual(remote.project, "other-project-123")
682682

683+
def test_referenced_property_graphs(self):
684+
from google.cloud.bigquery.table import PropertyGraphReference
685+
686+
ref_pg_resource = [
687+
{
688+
"projectId": self.PROJECT,
689+
"datasetId": "dataset",
690+
"propertyGraphId": "pg1",
691+
},
692+
{
693+
"projectId": self.PROJECT,
694+
"datasetId": "dataset",
695+
"propertyGraphId": "pg2",
696+
},
697+
{
698+
"projectId": "other-project-123",
699+
"datasetId": "other-dataset",
700+
"propertyGraphId": "other-pg",
701+
},
702+
]
703+
client = _make_client(project=self.PROJECT)
704+
job = self._make_one(self.JOB_ID, self.QUERY, client)
705+
self.assertEqual(job.referenced_property_graphs, [])
706+
707+
statistics = job._properties["statistics"] = {}
708+
self.assertEqual(job.referenced_property_graphs, [])
709+
710+
query_stats = statistics["query"] = {}
711+
self.assertEqual(job.referenced_property_graphs, [])
712+
713+
query_stats["referencedPropertyGraphs"] = ref_pg_resource
714+
715+
pg1, pg2, remote = job.referenced_property_graphs
716+
717+
self.assertIsInstance(pg1, PropertyGraphReference)
718+
self.assertEqual(pg1.property_graph_id, "pg1")
719+
self.assertEqual(pg1.dataset_id, "dataset")
720+
self.assertEqual(pg1.project, self.PROJECT)
721+
722+
self.assertIsInstance(pg2, PropertyGraphReference)
723+
self.assertEqual(pg2.property_graph_id, "pg2")
724+
self.assertEqual(pg2.dataset_id, "dataset")
725+
self.assertEqual(pg2.project, self.PROJECT)
726+
727+
self.assertIsInstance(remote, PropertyGraphReference)
728+
self.assertEqual(remote.property_graph_id, "other-pg")
729+
self.assertEqual(remote.dataset_id, "other-dataset")
730+
self.assertEqual(remote.project, "other-project-123")
731+
683732
def test_timeline(self):
684733
timeline_resource = [
685734
{

0 commit comments

Comments
 (0)