Skip to content

Commit fa9f9ca

Browse files
aravinsivatswast
aravinsiva
and
tswast
authored
feat: add instrumentation to list methods (#239)
* testing first trace export * instrumention client.py * instrumenting job.py and adding documentation * reconfiguring imports * quick cleanup of unused variable * adding more attributes in module and limiting complexity of instrumentation * adding tests, nox and correct attribute additions in client & job * adding tests, nox and correct attribute additions in client & job (left out of last commit) * linting * reformatting noxfile.[y * addressing suggested changes * adding suggested changes * removing print statements * setting same version across all OT [ackages and other reccommended changes * suggested changes * fixing packages issue in nox and updating documentation * fixing module install issue * restructuring design for testing adding first layer of tests (some still failing) * adding reamining client tests and all job tests * fixing linting issues * fixing trace not defined issue * fixing lint issues * fixing documentation issues and python2 testing issue * linting and fixing coverage issues * adding suggested changes * linting * adding Shawn's suggested changes * fixing _default_span_attribute_bug * reverting uneccesxsary changes * adding more tests for all job_ref parameters * removing dependecny, ordering imports and other changes * addressing Shawn concerns * adding test and suggested changes * adding opentelemetry to setup.py and other suggested changes * adding reasoning for not adding to [all] * linting * adding nested functions * adding test * adding Tim suggested changes * adding full tests * linting and fixing missing test Co-authored-by: Tim Swast <[email protected]>
1 parent 2874abf commit fa9f9ca

File tree

2 files changed

+162
-23
lines changed

2 files changed

+162
-23
lines changed

‎google/cloud/bigquery/client.py‎

Lines changed: 57 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -289,9 +289,17 @@ def list_projects(
289289
Iterator of :class:`~google.cloud.bigquery.client.Project`
290290
accessible to the current client.
291291
"""
292+
span_attributes = {"path": "/projects"}
293+
294+
def api_request(*args, **kwargs):
295+
with create_span(
296+
name="BigQuery.listProjects", attributes=span_attributes, client=self
297+
):
298+
return self._call_api(retry, *args, timeout=timeout, **kwargs)
299+
292300
return page_iterator.HTTPIterator(
293301
client=self,
294-
api_request=functools.partial(self._call_api, retry, timeout=timeout),
302+
api_request=api_request,
295303
path="/projects",
296304
item_to_value=_item_to_project,
297305
items_key="projects",
@@ -353,9 +361,18 @@ def list_datasets(
353361
# and converting it into a string here.
354362
extra_params["filter"] = filter
355363
path = "/projects/%s/datasets" % (project,)
364+
365+
span_attributes = {"path": path}
366+
367+
def api_request(*args, **kwargs):
368+
with create_span(
369+
name="BigQuery.listDatasets", attributes=span_attributes, client=self
370+
):
371+
return self._call_api(retry, *args, timeout=timeout, **kwargs)
372+
356373
return page_iterator.HTTPIterator(
357374
client=self,
358-
api_request=functools.partial(self._call_api, retry, timeout=timeout),
375+
api_request=api_request,
359376
path=path,
360377
item_to_value=_item_to_dataset,
361378
items_key="datasets",
@@ -1067,9 +1084,17 @@ def list_models(
10671084
raise TypeError("dataset must be a Dataset, DatasetReference, or string")
10681085

10691086
path = "%s/models" % dataset.path
1087+
span_attributes = {"path": path}
1088+
1089+
def api_request(*args, **kwargs):
1090+
with create_span(
1091+
name="BigQuery.listModels", attributes=span_attributes, client=self
1092+
):
1093+
return self._call_api(retry, *args, timeout=timeout, **kwargs)
1094+
10701095
result = page_iterator.HTTPIterator(
10711096
client=self,
1072-
api_request=functools.partial(self._call_api, retry, timeout=timeout),
1097+
api_request=api_request,
10731098
path=path,
10741099
item_to_value=_item_to_model,
10751100
items_key="models",
@@ -1132,9 +1157,18 @@ def list_routines(
11321157
raise TypeError("dataset must be a Dataset, DatasetReference, or string")
11331158

11341159
path = "{}/routines".format(dataset.path)
1160+
1161+
span_attributes = {"path": path}
1162+
1163+
def api_request(*args, **kwargs):
1164+
with create_span(
1165+
name="BigQuery.listRoutines", attributes=span_attributes, client=self
1166+
):
1167+
return self._call_api(retry, *args, timeout=timeout, **kwargs)
1168+
11351169
result = page_iterator.HTTPIterator(
11361170
client=self,
1137-
api_request=functools.partial(self._call_api, retry, timeout=timeout),
1171+
api_request=api_request,
11381172
path=path,
11391173
item_to_value=_item_to_routine,
11401174
items_key="routines",
@@ -1197,9 +1231,17 @@ def list_tables(
11971231
raise TypeError("dataset must be a Dataset, DatasetReference, or string")
11981232

11991233
path = "%s/tables" % dataset.path
1234+
span_attributes = {"path": path}
1235+
1236+
def api_request(*args, **kwargs):
1237+
with create_span(
1238+
name="BigQuery.listTables", attributes=span_attributes, client=self
1239+
):
1240+
return self._call_api(retry, *args, timeout=timeout, **kwargs)
1241+
12001242
result = page_iterator.HTTPIterator(
12011243
client=self,
1202-
api_request=functools.partial(self._call_api, retry, timeout=timeout),
1244+
api_request=api_request,
12031245
path=path,
12041246
item_to_value=_item_to_table,
12051247
items_key="tables",
@@ -1765,9 +1807,18 @@ def list_jobs(
17651807
project = self.project
17661808

17671809
path = "/projects/%s/jobs" % (project,)
1810+
1811+
span_attributes = {"path": path}
1812+
1813+
def api_request(*args, **kwargs):
1814+
with create_span(
1815+
name="BigQuery.listJobs", attributes=span_attributes, client=self
1816+
):
1817+
return self._call_api(retry, *args, timeout=timeout, **kwargs)
1818+
17681819
return page_iterator.HTTPIterator(
17691820
client=self,
1770-
api_request=functools.partial(self._call_api, retry, timeout=timeout),
1821+
api_request=api_request,
17711822
path=path,
17721823
item_to_value=_item_to_job,
17731824
items_key="jobs",

‎tests/unit/test_client.py‎

Lines changed: 105 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -425,9 +425,14 @@ def test_list_projects_defaults(self):
425425
creds = _make_credentials()
426426
client = self._make_one(PROJECT_1, creds)
427427
conn = client._connection = make_connection(DATA)
428-
429428
iterator = client.list_projects()
430-
page = six.next(iterator.pages)
429+
430+
with mock.patch(
431+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
432+
) as final_attributes:
433+
page = six.next(iterator.pages)
434+
435+
final_attributes.assert_called_once_with({"path": "/projects"}, client, None)
431436
projects = list(page)
432437
token = iterator.next_page_token
433438

@@ -455,7 +460,13 @@ def test_list_projects_w_timeout(self):
455460
conn = client._connection = make_connection(DATA)
456461

457462
iterator = client.list_projects(timeout=7.5)
458-
six.next(iterator.pages)
463+
464+
with mock.patch(
465+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
466+
) as final_attributes:
467+
six.next(iterator.pages)
468+
469+
final_attributes.assert_called_once_with({"path": "/projects"}, client, None)
459470

460471
conn.api_request.assert_called_once_with(
461472
method="GET", path="/projects", query_params={}, timeout=7.5
@@ -469,7 +480,13 @@ def test_list_projects_explicit_response_missing_projects_key(self):
469480
conn = client._connection = make_connection(DATA)
470481

471482
iterator = client.list_projects(max_results=3, page_token=TOKEN)
472-
page = six.next(iterator.pages)
483+
484+
with mock.patch(
485+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
486+
) as final_attributes:
487+
page = six.next(iterator.pages)
488+
489+
final_attributes.assert_called_once_with({"path": "/projects"}, client, None)
473490
projects = list(page)
474491
token = iterator.next_page_token
475492

@@ -518,7 +535,12 @@ def test_list_datasets_defaults(self):
518535
conn = client._connection = make_connection(DATA)
519536

520537
iterator = client.list_datasets()
521-
page = six.next(iterator.pages)
538+
with mock.patch(
539+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
540+
) as final_attributes:
541+
page = six.next(iterator.pages)
542+
543+
final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
522544
datasets = list(page)
523545
token = iterator.next_page_token
524546

@@ -538,7 +560,14 @@ def test_list_datasets_w_project_and_timeout(self):
538560
client = self._make_one(self.PROJECT, creds)
539561
conn = client._connection = make_connection({})
540562

541-
list(client.list_datasets(project="other-project", timeout=7.5))
563+
with mock.patch(
564+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
565+
) as final_attributes:
566+
list(client.list_datasets(project="other-project", timeout=7.5))
567+
568+
final_attributes.assert_called_once_with(
569+
{"path": "/projects/other-project/datasets"}, client, None
570+
)
542571

543572
conn.api_request.assert_called_once_with(
544573
method="GET",
@@ -559,7 +588,12 @@ def test_list_datasets_explicit_response_missing_datasets_key(self):
559588
iterator = client.list_datasets(
560589
include_all=True, filter=FILTER, max_results=3, page_token=TOKEN
561590
)
562-
page = six.next(iterator.pages)
591+
with mock.patch(
592+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
593+
) as final_attributes:
594+
page = six.next(iterator.pages)
595+
596+
final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
563597
datasets = list(page)
564598
token = iterator.next_page_token
565599

@@ -2838,7 +2872,12 @@ def test_list_tables_empty_w_timeout(self):
28382872
dataset = DatasetReference(self.PROJECT, self.DS_ID)
28392873
iterator = client.list_tables(dataset, timeout=7.5)
28402874
self.assertIs(iterator.dataset, dataset)
2841-
page = six.next(iterator.pages)
2875+
with mock.patch(
2876+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
2877+
) as final_attributes:
2878+
page = six.next(iterator.pages)
2879+
2880+
final_attributes.assert_called_once_with({"path": path}, client, None)
28422881
tables = list(page)
28432882
token = iterator.next_page_token
28442883

@@ -2856,7 +2895,12 @@ def test_list_models_empty_w_timeout(self):
28562895

28572896
dataset_id = "{}.{}".format(self.PROJECT, self.DS_ID)
28582897
iterator = client.list_models(dataset_id, timeout=7.5)
2859-
page = six.next(iterator.pages)
2898+
with mock.patch(
2899+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
2900+
) as final_attributes:
2901+
page = six.next(iterator.pages)
2902+
2903+
final_attributes.assert_called_once_with({"path": path}, client, None)
28602904
models = list(page)
28612905
token = iterator.next_page_token
28622906

@@ -2900,7 +2944,12 @@ def test_list_models_defaults(self):
29002944

29012945
iterator = client.list_models(dataset)
29022946
self.assertIs(iterator.dataset, dataset)
2903-
page = six.next(iterator.pages)
2947+
with mock.patch(
2948+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
2949+
) as final_attributes:
2950+
page = six.next(iterator.pages)
2951+
2952+
final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
29042953
models = list(page)
29052954
token = iterator.next_page_token
29062955

@@ -2926,7 +2975,16 @@ def test_list_routines_empty_w_timeout(self):
29262975
conn = client._connection = make_connection({})
29272976

29282977
iterator = client.list_routines("test-routines.test_routines", timeout=7.5)
2929-
page = six.next(iterator.pages)
2978+
with mock.patch(
2979+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
2980+
) as final_attributes:
2981+
page = six.next(iterator.pages)
2982+
2983+
final_attributes.assert_called_once_with(
2984+
{"path": "/projects/test-routines/datasets/test_routines/routines"},
2985+
client,
2986+
None,
2987+
)
29302988
routines = list(page)
29312989
token = iterator.next_page_token
29322990

@@ -2975,7 +3033,12 @@ def test_list_routines_defaults(self):
29753033

29763034
iterator = client.list_routines(dataset)
29773035
self.assertIs(iterator.dataset, dataset)
2978-
page = six.next(iterator.pages)
3036+
with mock.patch(
3037+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
3038+
) as final_attributes:
3039+
page = six.next(iterator.pages)
3040+
3041+
final_attributes.assert_called_once_with({"path": path}, client, None)
29793042
routines = list(page)
29803043
actual_token = iterator.next_page_token
29813044

@@ -3039,7 +3102,12 @@ def test_list_tables_defaults(self):
30393102

30403103
iterator = client.list_tables(dataset)
30413104
self.assertIs(iterator.dataset, dataset)
3042-
page = six.next(iterator.pages)
3105+
with mock.patch(
3106+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
3107+
) as final_attributes:
3108+
page = six.next(iterator.pages)
3109+
3110+
final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
30433111
tables = list(page)
30443112
token = iterator.next_page_token
30453113

@@ -3098,7 +3166,12 @@ def test_list_tables_explicit(self):
30983166
page_token=TOKEN,
30993167
)
31003168
self.assertEqual(iterator.dataset, dataset)
3101-
page = six.next(iterator.pages)
3169+
with mock.patch(
3170+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
3171+
) as final_attributes:
3172+
page = six.next(iterator.pages)
3173+
3174+
final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
31023175
tables = list(page)
31033176
token = iterator.next_page_token
31043177

@@ -3921,7 +3994,12 @@ def test_list_jobs_defaults(self):
39213994
conn = client._connection = make_connection(DATA)
39223995

39233996
iterator = client.list_jobs()
3924-
page = six.next(iterator.pages)
3997+
with mock.patch(
3998+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
3999+
) as final_attributes:
4000+
page = six.next(iterator.pages)
4001+
4002+
final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
39254003
jobs = list(page)
39264004
token = iterator.next_page_token
39274005

@@ -3966,7 +4044,12 @@ def test_list_jobs_load_job_wo_sourceUris(self):
39664044
conn = client._connection = make_connection(DATA)
39674045

39684046
iterator = client.list_jobs()
3969-
page = six.next(iterator.pages)
4047+
with mock.patch(
4048+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
4049+
) as final_attributes:
4050+
page = six.next(iterator.pages)
4051+
4052+
final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
39704053
jobs = list(page)
39714054
token = iterator.next_page_token
39724055

@@ -3995,7 +4078,12 @@ def test_list_jobs_explicit_missing(self):
39954078
iterator = client.list_jobs(
39964079
max_results=1000, page_token=TOKEN, all_users=True, state_filter="done"
39974080
)
3998-
page = six.next(iterator.pages)
4081+
with mock.patch(
4082+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
4083+
) as final_attributes:
4084+
page = six.next(iterator.pages)
4085+
4086+
final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
39994087
jobs = list(page)
40004088
token = iterator.next_page_token
40014089

0 commit comments

Comments
 (0)