From fa9f9ca491c3f9954287102c567ec483aa6151d4 Mon Sep 17 00:00:00 2001 From: Aravin <34178459+aravinsiva@users.noreply.github.com> Date: Mon, 24 Aug 2020 14:31:18 -0400 Subject: [PATCH] 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 --- google/cloud/bigquery/client.py | 63 +++++++++++++++-- tests/unit/test_client.py | 122 +++++++++++++++++++++++++++----- 2 files changed, 162 insertions(+), 23 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index fbbfda051..e7f8c5c66 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/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", diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 271640dd5..01bb1f2e1 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -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