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]>

Author: aravinsiva

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",

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