Skip to content
Permalink
Browse files
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 <swast@google.com>
  • Loading branch information
aravinsiva and tswast committed Aug 24, 2020
1 parent 2874abf commit fa9f9ca491c3f9954287102c567ec483aa6151d4
Showing with 162 additions and 23 deletions.
  1. +57 −6 google/cloud/bigquery/client.py
  2. +105 −17 tests/unit/test_client.py
@@ -289,9 +289,17 @@ def list_projects(
Iterator of :class:`~google.cloud.bigquery.client.Project`
accessible to the current client.
"""
span_attributes = {"path": "/projects"}

def api_request(*args, **kwargs):
with create_span(
name="BigQuery.listProjects", attributes=span_attributes, client=self
):
return self._call_api(retry, *args, timeout=timeout, **kwargs)

return page_iterator.HTTPIterator(
client=self,
api_request=functools.partial(self._call_api, retry, timeout=timeout),
api_request=api_request,
path="/projects",
item_to_value=_item_to_project,
items_key="projects",
@@ -353,9 +361,18 @@ def list_datasets(
# and converting it into a string here.
extra_params["filter"] = filter
path = "/projects/%s/datasets" % (project,)

span_attributes = {"path": path}

def api_request(*args, **kwargs):
with create_span(
name="BigQuery.listDatasets", attributes=span_attributes, client=self
):
return self._call_api(retry, *args, timeout=timeout, **kwargs)

return page_iterator.HTTPIterator(
client=self,
api_request=functools.partial(self._call_api, retry, timeout=timeout),
api_request=api_request,
path=path,
item_to_value=_item_to_dataset,
items_key="datasets",
@@ -1067,9 +1084,17 @@ def list_models(
raise TypeError("dataset must be a Dataset, DatasetReference, or string")

path = "%s/models" % dataset.path
span_attributes = {"path": path}

def api_request(*args, **kwargs):
with create_span(
name="BigQuery.listModels", attributes=span_attributes, client=self
):
return self._call_api(retry, *args, timeout=timeout, **kwargs)

result = page_iterator.HTTPIterator(
client=self,
api_request=functools.partial(self._call_api, retry, timeout=timeout),
api_request=api_request,
path=path,
item_to_value=_item_to_model,
items_key="models",
@@ -1132,9 +1157,18 @@ def list_routines(
raise TypeError("dataset must be a Dataset, DatasetReference, or string")

path = "{}/routines".format(dataset.path)

span_attributes = {"path": path}

def api_request(*args, **kwargs):
with create_span(
name="BigQuery.listRoutines", attributes=span_attributes, client=self
):
return self._call_api(retry, *args, timeout=timeout, **kwargs)

result = page_iterator.HTTPIterator(
client=self,
api_request=functools.partial(self._call_api, retry, timeout=timeout),
api_request=api_request,
path=path,
item_to_value=_item_to_routine,
items_key="routines",
@@ -1197,9 +1231,17 @@ def list_tables(
raise TypeError("dataset must be a Dataset, DatasetReference, or string")

path = "%s/tables" % dataset.path
span_attributes = {"path": path}

def api_request(*args, **kwargs):
with create_span(
name="BigQuery.listTables", attributes=span_attributes, client=self
):
return self._call_api(retry, *args, timeout=timeout, **kwargs)

result = page_iterator.HTTPIterator(
client=self,
api_request=functools.partial(self._call_api, retry, timeout=timeout),
api_request=api_request,
path=path,
item_to_value=_item_to_table,
items_key="tables",
@@ -1765,9 +1807,18 @@ def list_jobs(
project = self.project

path = "/projects/%s/jobs" % (project,)

span_attributes = {"path": path}

def api_request(*args, **kwargs):
with create_span(
name="BigQuery.listJobs", attributes=span_attributes, client=self
):
return self._call_api(retry, *args, timeout=timeout, **kwargs)

return page_iterator.HTTPIterator(
client=self,
api_request=functools.partial(self._call_api, retry, timeout=timeout),
api_request=api_request,
path=path,
item_to_value=_item_to_job,
items_key="jobs",
@@ -425,9 +425,14 @@ def test_list_projects_defaults(self):
creds = _make_credentials()
client = self._make_one(PROJECT_1, creds)
conn = client._connection = make_connection(DATA)

iterator = client.list_projects()
page = six.next(iterator.pages)

with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/projects"}, client, None)
projects = list(page)
token = iterator.next_page_token

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

iterator = client.list_projects(timeout=7.5)
six.next(iterator.pages)

with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/projects"}, client, None)

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

iterator = client.list_projects(max_results=3, page_token=TOKEN)
page = six.next(iterator.pages)

with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/projects"}, client, None)
projects = list(page)
token = iterator.next_page_token

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

iterator = client.list_datasets()
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
datasets = list(page)
token = iterator.next_page_token

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

list(client.list_datasets(project="other-project", timeout=7.5))
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
list(client.list_datasets(project="other-project", timeout=7.5))

final_attributes.assert_called_once_with(
{"path": "/projects/other-project/datasets"}, client, None
)

conn.api_request.assert_called_once_with(
method="GET",
@@ -559,7 +588,12 @@ def test_list_datasets_explicit_response_missing_datasets_key(self):
iterator = client.list_datasets(
include_all=True, filter=FILTER, max_results=3, page_token=TOKEN
)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
datasets = list(page)
token = iterator.next_page_token

@@ -2838,7 +2872,12 @@ def test_list_tables_empty_w_timeout(self):
dataset = DatasetReference(self.PROJECT, self.DS_ID)
iterator = client.list_tables(dataset, timeout=7.5)
self.assertIs(iterator.dataset, dataset)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": path}, client, None)
tables = list(page)
token = iterator.next_page_token

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

dataset_id = "{}.{}".format(self.PROJECT, self.DS_ID)
iterator = client.list_models(dataset_id, timeout=7.5)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": path}, client, None)
models = list(page)
token = iterator.next_page_token

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

iterator = client.list_models(dataset)
self.assertIs(iterator.dataset, dataset)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
models = list(page)
token = iterator.next_page_token

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

iterator = client.list_routines("test-routines.test_routines", timeout=7.5)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with(
{"path": "/projects/test-routines/datasets/test_routines/routines"},
client,
None,
)
routines = list(page)
token = iterator.next_page_token

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

iterator = client.list_routines(dataset)
self.assertIs(iterator.dataset, dataset)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": path}, client, None)
routines = list(page)
actual_token = iterator.next_page_token

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

iterator = client.list_tables(dataset)
self.assertIs(iterator.dataset, dataset)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
tables = list(page)
token = iterator.next_page_token

@@ -3098,7 +3166,12 @@ def test_list_tables_explicit(self):
page_token=TOKEN,
)
self.assertEqual(iterator.dataset, dataset)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
tables = list(page)
token = iterator.next_page_token

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

iterator = client.list_jobs()
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
jobs = list(page)
token = iterator.next_page_token

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

iterator = client.list_jobs()
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
jobs = list(page)
token = iterator.next_page_token

@@ -3995,7 +4078,12 @@ def test_list_jobs_explicit_missing(self):
iterator = client.list_jobs(
max_results=1000, page_token=TOKEN, all_users=True, state_filter="done"
)
page = six.next(iterator.pages)
with mock.patch(
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
) as final_attributes:
page = six.next(iterator.pages)

final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None)
jobs = list(page)
token = iterator.next_page_token

0 comments on commit fa9f9ca

Please sign in to comment.