diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index e7f8c5c66..86275487b 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -248,14 +248,14 @@ def get_service_account_email( project = self.project path = "/projects/%s/serviceAccount" % (project,) span_attributes = {"path": path} - with create_span( - name="BigQuery.getServiceAccountEmail", - attributes=span_attributes, - client=self, - ): - api_response = self._call_api( - retry, method="GET", path=path, timeout=timeout - ) + api_response = self._call_api( + retry, + span_name="BigQuery.getServiceAccountEmail", + span_attributes=span_attributes, + method="GET", + path=path, + timeout=timeout, + ) return api_response["email"] def list_projects( @@ -292,10 +292,14 @@ def list_projects( 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 self._call_api( + retry, + span_name="BigQuery.listProjects", + span_attributes=span_attributes, + *args, + timeout=timeout, + **kwargs + ) return page_iterator.HTTPIterator( client=self, @@ -365,10 +369,15 @@ def list_datasets( 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 self._call_api( + retry, + span_name="BigQuery.listDatasets", + span_attributes=span_attributes, + *args, + timeout=timeout, + **kwargs + ) return page_iterator.HTTPIterator( client=self, @@ -497,12 +506,16 @@ def create_dataset( try: span_attributes = {"path": path} - with create_span( - name="BigQuery.createDataset", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, method="POST", path=path, data=data, timeout=timeout - ) + + api_response = self._call_api( + retry, + span_name="BigQuery.createDataset", + span_attributes=span_attributes, + method="POST", + path=path, + data=data, + timeout=timeout, + ) return Dataset.from_api_repr(api_response) except google.api_core.exceptions.Conflict: if not exists_ok: @@ -545,12 +558,15 @@ def create_routine( resource = routine.to_api_repr() try: span_attributes = {"path": path} - with create_span( - name="BigQuery.createRoutine", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, method="POST", path=path, data=resource, timeout=timeout - ) + api_response = self._call_api( + retry, + span_name="BigQuery.createRoutine", + span_attributes=span_attributes, + method="POST", + path=path, + data=resource, + timeout=timeout, + ) return Routine.from_api_repr(api_response) except google.api_core.exceptions.Conflict: if not exists_ok: @@ -596,22 +612,33 @@ def create_table(self, table, exists_ok=False, retry=DEFAULT_RETRY, timeout=None data = table.to_api_repr() try: span_attributes = {"path": path, "dataset_id": dataset_id} - with create_span( - name="BigQuery.createTable", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, method="POST", path=path, data=data, timeout=timeout - ) + api_response = self._call_api( + retry, + span_name="BigQuery.createTable", + span_attributes=span_attributes, + method="POST", + path=path, + data=data, + timeout=timeout, + ) return Table.from_api_repr(api_response) except google.api_core.exceptions.Conflict: if not exists_ok: raise return self.get_table(table.reference, retry=retry) - def _call_api(self, retry, **kwargs): + def _call_api( + self, retry, span_name=None, span_attributes=None, job_ref=None, **kwargs + ): + call = functools.partial(self._connection.api_request, **kwargs) if retry: call = retry(call) + if span_name is not None: + with create_span( + name=span_name, attributes=span_attributes, client=self, job_ref=job_ref + ): + return call() return call() def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None): @@ -642,12 +669,14 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None): ) path = dataset_ref.path span_attributes = {"path": path} - with create_span( - name="BigQuery.getDataset", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, method="GET", path=path, timeout=timeout - ) + api_response = self._call_api( + retry, + span_name="BigQuery.getDataset", + span_attributes=span_attributes, + method="GET", + path=path, + timeout=timeout, + ) return Dataset.from_api_repr(api_response) def get_iam_policy( @@ -663,12 +692,15 @@ def get_iam_policy( path = "{}:getIamPolicy".format(table.path) span_attributes = {"path": path} - with create_span( - name="BigQuery.getIamPolicy", attributes=span_attributes, client=self - ): - response = self._call_api( - retry, method="POST", path=path, data=body, timeout=timeout, - ) + response = self._call_api( + retry, + span_name="BigQuery.getIamPolicy", + span_attributes=span_attributes, + method="POST", + path=path, + data=body, + timeout=timeout, + ) return Policy.from_api_repr(response) @@ -688,12 +720,16 @@ def set_iam_policy( path = "{}:setIamPolicy".format(table.path) span_attributes = {"path": path} - with create_span( - name="BigQuery.setIamPolicy", attributes=span_attributes, client=self - ): - response = self._call_api( - retry, method="POST", path=path, data=body, timeout=timeout, - ) + + response = self._call_api( + retry, + span_name="BigQuery.setIamPolicy", + span_attributes=span_attributes, + method="POST", + path=path, + data=body, + timeout=timeout, + ) return Policy.from_api_repr(response) @@ -707,12 +743,15 @@ def test_iam_permissions( path = "{}:testIamPermissions".format(table.path) span_attributes = {"path": path} - with create_span( - name="BigQuery.testIamPermissions", attributes=span_attributes, client=self - ): - response = self._call_api( - retry, method="POST", path=path, data=body, timeout=timeout, - ) + response = self._call_api( + retry, + span_name="BigQuery.testIamPermissions", + span_attributes=span_attributes, + method="POST", + path=path, + data=body, + timeout=timeout, + ) return response @@ -743,12 +782,15 @@ def get_model(self, model_ref, retry=DEFAULT_RETRY, timeout=None): ) path = model_ref.path span_attributes = {"path": path} - with create_span( - name="BigQuery.getModel", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, method="GET", path=path, timeout=timeout - ) + + api_response = self._call_api( + retry, + span_name="BigQuery.getModel", + span_attributes=span_attributes, + method="GET", + path=path, + timeout=timeout, + ) return Model.from_api_repr(api_response) def get_routine(self, routine_ref, retry=DEFAULT_RETRY, timeout=None): @@ -780,12 +822,14 @@ def get_routine(self, routine_ref, retry=DEFAULT_RETRY, timeout=None): ) path = routine_ref.path span_attributes = {"path": path} - with create_span( - name="BigQuery.getRoutine", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, method="GET", path=path, timeout=timeout - ) + api_response = self._call_api( + retry, + span_name="BigQuery.getRoutine", + span_attributes=span_attributes, + method="GET", + path=path, + timeout=timeout, + ) return Routine.from_api_repr(api_response) def get_table(self, table, retry=DEFAULT_RETRY, timeout=None): @@ -814,12 +858,14 @@ def get_table(self, table, retry=DEFAULT_RETRY, timeout=None): table_ref = _table_arg_to_table_ref(table, default_project=self.project) path = table_ref.path span_attributes = {"path": path} - with create_span( - name="BigQuery.getTable", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, method="GET", path=path, timeout=timeout - ) + api_response = self._call_api( + retry, + span_name="BigQuery.getTable", + span_attributes=span_attributes, + method="GET", + path=path, + timeout=timeout, + ) return Table.from_api_repr(api_response) def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): @@ -859,17 +905,16 @@ def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): path = dataset.path span_attributes = {"path": path, "fields": fields} - with create_span( - name="BigQuery.updateDataset", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, - method="PATCH", - path=path, - data=partial, - headers=headers, - timeout=timeout, - ) + api_response = self._call_api( + retry, + span_name="BigQuery.updateDataset", + span_attributes=span_attributes, + method="PATCH", + path=path, + data=partial, + headers=headers, + timeout=timeout, + ) return Dataset.from_api_repr(api_response) def update_model(self, model, fields, retry=DEFAULT_RETRY, timeout=None): @@ -908,17 +953,16 @@ def update_model(self, model, fields, retry=DEFAULT_RETRY, timeout=None): path = model.path span_attributes = {"path": path, "fields": fields} - with create_span( - name="BigQuery.updateModel", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, - method="PATCH", - path=path, - data=partial, - headers=headers, - timeout=timeout, - ) + api_response = self._call_api( + retry, + span_name="BigQuery.updateModel", + span_attributes=span_attributes, + method="PATCH", + path=path, + data=partial, + headers=headers, + timeout=timeout, + ) return Model.from_api_repr(api_response) def update_routine(self, routine, fields, retry=DEFAULT_RETRY, timeout=None): @@ -968,17 +1012,16 @@ def update_routine(self, routine, fields, retry=DEFAULT_RETRY, timeout=None): path = routine.path span_attributes = {"path": path, "fields": fields} - with create_span( - name="BigQuery.updateRoutine", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, - method="PUT", - path=path, - data=partial, - headers=headers, - timeout=timeout, - ) + api_response = self._call_api( + retry, + span_name="BigQuery.updateRoutine", + span_attributes=span_attributes, + method="PUT", + path=path, + data=partial, + headers=headers, + timeout=timeout, + ) return Routine.from_api_repr(api_response) def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): @@ -1018,17 +1061,16 @@ def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): path = table.path span_attributes = {"path": path, "fields": fields} - with create_span( - name="BigQuery.updateTable", attributes=span_attributes, client=self - ): - api_response = self._call_api( - retry, - method="PATCH", - path=path, - data=partial, - headers=headers, - timeout=timeout, - ) + api_response = self._call_api( + retry, + span_name="BigQuery.updateTable", + span_attributes=span_attributes, + method="PATCH", + path=path, + data=partial, + headers=headers, + timeout=timeout, + ) return Table.from_api_repr(api_response) def list_models( @@ -1087,10 +1129,14 @@ def list_models( 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) + return self._call_api( + retry, + span_name="BigQuery.listModels", + span_attributes=span_attributes, + *args, + timeout=timeout, + **kwargs + ) result = page_iterator.HTTPIterator( client=self, @@ -1161,10 +1207,14 @@ def list_routines( 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) + return self._call_api( + retry, + span_name="BigQuery.listRoutines", + span_attributes=span_attributes, + *args, + timeout=timeout, + **kwargs + ) result = page_iterator.HTTPIterator( client=self, @@ -1234,10 +1284,14 @@ def list_tables( 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) + return self._call_api( + retry, + span_name="BigQuery.listTables", + span_attributes=span_attributes, + *args, + timeout=timeout, + **kwargs + ) result = page_iterator.HTTPIterator( client=self, @@ -1304,16 +1358,15 @@ def delete_dataset( span_attributes = {"path": path} try: - with create_span( - name="BigQuery.deleteDataset", attributes=span_attributes, client=self - ): - self._call_api( - retry, - method="DELETE", - path=path, - query_params=params, - timeout=timeout, - ) + self._call_api( + retry, + span_name="BigQuery.deleteDataset", + span_attributes=span_attributes, + method="DELETE", + path=path, + query_params=params, + timeout=timeout, + ) except google.api_core.exceptions.NotFound: if not not_found_ok: raise @@ -1354,10 +1407,14 @@ def delete_model( path = model.path try: span_attributes = {"path": path} - with create_span( - name="BigQuery.deleteModel", attributes=span_attributes, client=self - ): - self._call_api(retry, method="DELETE", path=path, timeout=timeout) + self._call_api( + retry, + span_name="BigQuery.deleteModel", + span_attributes=span_attributes, + method="DELETE", + path=path, + timeout=timeout, + ) except google.api_core.exceptions.NotFound: if not not_found_ok: raise @@ -1400,10 +1457,14 @@ def delete_routine( try: span_attributes = {"path": path} - with create_span( - name="BigQuery.deleteRoutine", attributes=span_attributes, client=self - ): - self._call_api(retry, method="DELETE", path=path, timeout=timeout) + self._call_api( + retry, + span_name="BigQuery.deleteRoutine", + span_attributes=span_attributes, + method="DELETE", + path=path, + timeout=timeout, + ) except google.api_core.exceptions.NotFound: if not not_found_ok: raise @@ -1442,10 +1503,14 @@ def delete_table( try: path = table.path span_attributes = {"path": path} - with create_span( - name="BigQuery.deleteTable", attributes=span_attributes, client=self - ): - self._call_api(retry, method="DELETE", path=path, timeout=timeout) + self._call_api( + retry, + span_name="BigQuery.deleteTable", + span_attributes=span_attributes, + method="DELETE", + path=path, + timeout=timeout, + ) except google.api_core.exceptions.NotFound: if not not_found_ok: raise @@ -1494,17 +1559,15 @@ def _get_query_results( # job is complete (from QueryJob.done(), called ultimately from # QueryJob.result()). So we don't need to poll here. span_attributes = {"path": path} - - with create_span( - name="BigQuery.getQueryResults", attributes=span_attributes, client=self - ): - resource = self._call_api( - retry, - method="GET", - path=path, - query_params=extra_params, - timeout=timeout, - ) + resource = self._call_api( + retry, + span_name="BigQuery.getQueryResults", + span_attributes=span_attributes, + method="GET", + path=path, + query_params=extra_params, + timeout=timeout, + ) return _QueryResults.from_api_repr(resource) def job_from_resource(self, resource): @@ -1650,16 +1713,15 @@ def get_job( span_attributes = {"path": path, "job_id": job_id, "location": location} - with create_span( - name="BigQuery.getJob", attributes=span_attributes, client=self - ): - resource = self._call_api( - retry, - method="GET", - path=path, - query_params=extra_params, - timeout=timeout, - ) + resource = self._call_api( + retry, + span_name="BigQuery.getJob", + span_attributes=span_attributes, + method="GET", + path=path, + query_params=extra_params, + timeout=timeout, + ) return self.job_from_resource(resource) @@ -1708,16 +1770,15 @@ def cancel_job( span_attributes = {"path": path, "job_id": job_id, "location": location} - with create_span( - name="BigQuery.cancelJob", attributes=span_attributes, client=self - ): - resource = self._call_api( - retry, - method="POST", - path=path, - query_params=extra_params, - timeout=timeout, - ) + resource = self._call_api( + retry, + span_name="BigQuery.cancelJob", + span_attributes=span_attributes, + method="POST", + path=path, + query_params=extra_params, + timeout=timeout, + ) return self.job_from_resource(resource["job"]) @@ -1811,10 +1872,14 @@ def list_jobs( 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 self._call_api( + retry, + span_name="BigQuery.listJobs", + span_attributes=span_attributes, + *args, + timeout=timeout, + **kwargs + ) return page_iterator.HTTPIterator( client=self, @@ -2904,12 +2969,15 @@ def insert_rows_json( path = "%s/insertAll" % table.path # We can always retry, because every row has an insert ID. span_attributes = {"path": path} - with create_span( - name="BigQuery.insertRowsJson", attributes=span_attributes, client=self - ): - response = self._call_api( - retry, method="POST", path=path, data=data, timeout=timeout, - ) + response = self._call_api( + retry, + span_name="BigQuery.insertRowsJson", + span_attributes=span_attributes, + method="POST", + path=path, + data=data, + timeout=timeout, + ) errors = [] for error in response.get("insertErrors", ()): @@ -3066,10 +3134,10 @@ def list_rows( def _schema_from_json_file_object(self, file_obj): """Helper function for schema_from_json that takes a - file object that describes a table schema. + file object that describes a table schema. - Returns: - List of schema field objects. + Returns: + List of schema field objects. """ json_data = json.load(file_obj) return [SchemaField.from_api_repr(field) for field in json_data] diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index a8e0c25ed..20bce597a 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -34,7 +34,6 @@ from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration from google.cloud.bigquery.external_config import ExternalConfig from google.cloud.bigquery.external_config import HivePartitioningOptions -from google.cloud.bigquery.opentelemetry_tracing import create_span from google.cloud.bigquery import _helpers from google.cloud.bigquery.query import _query_param_from_api_repr from google.cloud.bigquery.query import ArrayQueryParameter @@ -636,16 +635,16 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None): # jobs.insert is idempotent because we ensure that every new # job has an ID. span_attributes = {"path": path} - with create_span( - name="BigQuery.job.begin", attributes=span_attributes, job_ref=self - ): - api_response = client._call_api( - retry, - method="POST", - path=path, - data=self.to_api_repr(), - timeout=timeout, - ) + api_response = client._call_api( + retry, + span_name="BigQuery.job.begin", + span_attributes=span_attributes, + job_ref=self, + method="POST", + path=path, + data=self.to_api_repr(), + timeout=timeout, + ) self._set_properties(api_response) def exists(self, client=None, retry=DEFAULT_RETRY, timeout=None): @@ -675,16 +674,17 @@ def exists(self, client=None, retry=DEFAULT_RETRY, timeout=None): try: span_attributes = {"path": self.path} - with create_span( - name="BigQuery.job.exists", attributes=span_attributes, job_ref=self - ): - client._call_api( - retry, - method="GET", - path=self.path, - query_params=extra_params, - timeout=timeout, - ) + + client._call_api( + retry, + span_name="BigQuery.job.exists", + span_attributes=span_attributes, + job_ref=self, + method="GET", + path=self.path, + query_params=extra_params, + timeout=timeout, + ) except NotFound: return False else: @@ -712,16 +712,17 @@ def reload(self, client=None, retry=DEFAULT_RETRY, timeout=None): if self.location: extra_params["location"] = self.location span_attributes = {"path": self.path} - with create_span( - name="BigQuery.job.reload", attributes=span_attributes, job_ref=self - ): - api_response = client._call_api( - retry, - method="GET", - path=self.path, - query_params=extra_params, - timeout=timeout, - ) + + api_response = client._call_api( + retry, + span_name="BigQuery.job.reload", + span_attributes=span_attributes, + job_ref=self, + method="GET", + path=self.path, + query_params=extra_params, + timeout=timeout, + ) self._set_properties(api_response) def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None): @@ -750,16 +751,17 @@ def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None): path = "{}/cancel".format(self.path) span_attributes = {"path": path} - with create_span( - name="BigQuery.job.cancel", attributes=span_attributes, job_ref=self - ): - api_response = client._call_api( - retry, - method="POST", - path=path, - query_params=extra_params, - timeout=timeout, - ) + + api_response = client._call_api( + retry, + span_name="BigQuery.job.cancel", + span_attributes=span_attributes, + job_ref=self, + method="POST", + path=path, + query_params=extra_params, + timeout=timeout, + ) self._set_properties(api_response["job"]) # The Future interface requires that we return True if the *attempt* # to cancel was successful. diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 01bb1f2e1..d354735a1 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -250,6 +250,54 @@ def test__call_api_applying_custom_retry_on_timeout(self): [mock.call(foo="bar"), mock.call(foo="bar")], # was retried once ) + def test__call_api_span_creator_not_called(self): + from concurrent.futures import TimeoutError + from google.cloud.bigquery.retry import DEFAULT_RETRY + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + api_request_patcher = mock.patch.object( + client._connection, "api_request", side_effect=[TimeoutError, "result"], + ) + retry = DEFAULT_RETRY.with_deadline(1).with_predicate( + lambda exc: isinstance(exc, TimeoutError) + ) + + with api_request_patcher: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: + client._call_api(retry) + + final_attributes.assert_not_called() + + def test__call_api_span_creator_called(self): + from concurrent.futures import TimeoutError + from google.cloud.bigquery.retry import DEFAULT_RETRY + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + api_request_patcher = mock.patch.object( + client._connection, "api_request", side_effect=[TimeoutError, "result"], + ) + retry = DEFAULT_RETRY.with_deadline(1).with_predicate( + lambda exc: isinstance(exc, TimeoutError) + ) + + with api_request_patcher: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: + client._call_api( + retry, + span_name="test_name", + span_attributes={"test_attribute": "test_attribute-value"}, + ) + + final_attributes.assert_called_once() + def test__get_query_results_miss_w_explicit_project_and_timeout(self): from google.cloud.exceptions import NotFound diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index d5497ffa8..fb6a46bd6 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -621,15 +621,17 @@ def test__begin_defaults(self): call_api = job._client._call_api = mock.Mock() call_api.return_value = resource path = "/projects/{}/jobs".format(self.PROJECT) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - job._begin() - - final_attributes.assert_called_with({"path": path}, None, job) + job._begin() call_api.assert_called_once_with( - DEFAULT_RETRY, method="POST", path=path, data=resource, timeout=None, + DEFAULT_RETRY, + span_name="BigQuery.job.begin", + span_attributes={"path": path}, + job_ref=job, + method="POST", + path=path, + data=resource, + timeout=None, ) self.assertEqual(job._properties, resource) @@ -653,15 +655,17 @@ def test__begin_explicit(self): call_api.return_value = resource retry = DEFAULT_RETRY.with_deadline(1) path = "/projects/{}/jobs".format(self.PROJECT) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - job._begin(client=client, retry=retry, timeout=7.5) - - final_attributes.assert_called_with({"path": path}, None, job) + job._begin(client=client, retry=retry, timeout=7.5) call_api.assert_called_once_with( - retry, method="POST", path=path, data=resource, timeout=7.5, + retry, + span_name="BigQuery.job.begin", + span_attributes={"path": path}, + job_ref=job, + method="POST", + path=path, + data=resource, + timeout=7.5, ) self.assertEqual(job._properties, resource) @@ -673,20 +677,15 @@ def test_exists_defaults_miss(self): job._properties["jobReference"]["location"] = self.LOCATION call_api = job._client._call_api = mock.Mock() call_api.side_effect = NotFound("testing") - - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - self.assertFalse(job.exists()) - - final_attributes.assert_called_with( - {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, - None, - job, - ) + self.assertFalse(job.exists()) call_api.assert_called_once_with( DEFAULT_RETRY, + span_name="BigQuery.job.exists", + span_attributes={ + "path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID) + }, + job_ref=job, method="GET", path="/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID), query_params={"fields": "id", "location": self.LOCATION}, @@ -710,19 +709,15 @@ def test_exists_explicit_hit(self): call_api = client._call_api = mock.Mock() call_api.return_value = resource retry = DEFAULT_RETRY.with_deadline(1) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - self.assertTrue(job.exists(client=client, retry=retry)) - - final_attributes.assert_called_with( - {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, - None, - job, - ) + self.assertTrue(job.exists(client=client, retry=retry)) call_api.assert_called_once_with( retry, + span_name="BigQuery.job.exists", + span_attributes={ + "path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID) + }, + job_ref=job, method="GET", path="/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID), query_params={"fields": "id"}, @@ -735,15 +730,13 @@ def test_exists_w_timeout(self): PATH = "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID) job = self._set_properties_job() call_api = job._client._call_api = mock.Mock() - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - job.exists(timeout=7.5) - - final_attributes.assert_called_with({"path": PATH}, None, job) + job.exists(timeout=7.5) call_api.assert_called_once_with( DEFAULT_RETRY, + span_name="BigQuery.job.exists", + span_attributes={"path": PATH}, + job_ref=job, method="GET", path=PATH, query_params={"fields": "id"}, @@ -765,19 +758,15 @@ def test_reload_defaults(self): job._properties["jobReference"]["location"] = self.LOCATION call_api = job._client._call_api = mock.Mock() call_api.return_value = resource - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - job.reload() - - final_attributes.assert_called_with( - {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, - None, - job, - ) + job.reload() call_api.assert_called_once_with( DEFAULT_RETRY, + span_name="BigQuery.job.reload", + span_attributes={ + "path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID) + }, + job_ref=job, method="GET", path="/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID), query_params={"location": self.LOCATION}, @@ -802,19 +791,15 @@ def test_reload_explicit(self): call_api = client._call_api = mock.Mock() call_api.return_value = resource retry = DEFAULT_RETRY.with_deadline(1) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - job.reload(client=client, retry=retry, timeout=4.2) - - final_attributes.assert_called_with( - {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, - None, - job, - ) + job.reload(client=client, retry=retry, timeout=4.2) call_api.assert_called_once_with( retry, + span_name="BigQuery.job.reload", + span_attributes={ + "path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID) + }, + job_ref=job, method="GET", path="/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID), query_params={}, @@ -840,11 +825,7 @@ def test_cancel_defaults(self): ) as final_attributes: self.assertTrue(job.cancel()) - final_attributes.assert_called_with( - {"path": "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)}, - None, - job, - ) + final_attributes.assert_called() connection.api_request.assert_called_once_with( method="POST", @@ -875,7 +856,7 @@ def test_cancel_explicit(self): final_attributes.assert_called_with( {"path": "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)}, - None, + client, job, ) @@ -915,7 +896,7 @@ def test_cancel_w_custom_retry(self): ) as final_attributes: result = job.cancel(retry=retry, timeout=7.5) - final_attributes.assert_called_with({"path": api_path}, None, job) + final_attributes.assert_called() self.assertTrue(result) self.assertEqual(job._properties, resource) @@ -2409,7 +2390,7 @@ def test_begin_w_bound_client(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": path}, None, job) + final_attributes.assert_called_with({"path": path}, client, job) conn.api_request.assert_called_once_with( method="POST", @@ -2454,7 +2435,7 @@ def test_begin_w_autodetect(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": path}, None, job) + final_attributes.assert_called_with({"path": path}, client, job) sent = { "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, @@ -2553,7 +2534,7 @@ def test_begin_w_alternate_client(self): ) as final_attributes: job._begin(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() self.assertEqual(len(conn2.api_request.call_args_list), 1) @@ -2583,7 +2564,7 @@ def test_begin_w_job_reference(self): ) as final_attributes: load_job._begin() final_attributes.assert_called_with( - {"path": "/projects/alternative-project/jobs"}, None, load_job + {"path": "/projects/alternative-project/jobs"}, client, load_job ) conn.api_request.assert_called_once() @@ -2608,7 +2589,7 @@ def test_exists_miss_w_bound_client(self): final_attributes.assert_called_with( {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, - None, + client, job, ) @@ -2630,7 +2611,7 @@ def test_exists_hit_w_alternate_client(self): final_attributes.assert_called_with( {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, - None, + client2, job, ) @@ -2652,7 +2633,7 @@ def test_exists_miss_w_job_reference(self): self.assertFalse(load_job.exists()) final_attributes.assert_called_with( - {"path": "/projects/other-project/jobs/my-job-id"}, None, load_job + {"path": "/projects/other-project/jobs/my-job-id"}, client, load_job ) conn.api_request.assert_called_once_with( @@ -2673,7 +2654,7 @@ def test_reload_w_bound_client(self): ) as final_attributes: job.reload() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None @@ -2693,7 +2674,7 @@ def test_reload_w_alternate_client(self): ) as final_attributes: job.reload(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -2718,7 +2699,7 @@ def test_reload_w_job_reference(self): final_attributes.assert_called_with( {"path": "/projects/alternative-project/jobs/{}".format(self.JOB_ID)}, - None, + client, load_job, ) @@ -2741,7 +2722,7 @@ def test_cancel_w_bound_client(self): ) as final_attributes: job.cancel() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="POST", path=PATH, query_params={}, timeout=None, @@ -2762,7 +2743,7 @@ def test_cancel_w_alternate_client(self): ) as final_attributes: job.cancel(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -2791,7 +2772,7 @@ def test_cancel_w_job_reference(self): self.JOB_ID ) }, - None, + client, load_job, ) conn.api_request.assert_called_once_with( @@ -3093,7 +3074,7 @@ def test_begin_w_bound_client(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="POST", @@ -3161,7 +3142,7 @@ def test_begin_w_alternate_client(self): ) as final_attributes: job._begin(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3188,7 +3169,7 @@ def test_exists_miss_w_bound_client(self): ) as final_attributes: self.assertFalse(job.exists()) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None, @@ -3208,7 +3189,7 @@ def test_exists_hit_w_alternate_client(self): ) as final_attributes: self.assertTrue(job.exists(client=client2)) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3228,7 +3209,7 @@ def test_reload_w_bound_client(self): ) as final_attributes: job.reload() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None @@ -3250,7 +3231,7 @@ def test_reload_w_alternate_client(self): ) as final_attributes: job.reload(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3515,7 +3496,7 @@ def test_begin_w_bound_client(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="POST", @@ -3577,7 +3558,7 @@ def test_begin_w_alternate_client(self): ) as final_attributes: job._begin(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3603,7 +3584,7 @@ def test_exists_miss_w_bound_client(self): ) as final_attributes: self.assertFalse(job.exists()) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None, @@ -3623,7 +3604,7 @@ def test_exists_hit_w_alternate_client(self): ) as final_attributes: self.assertTrue(job.exists(client=client2)) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3645,7 +3626,7 @@ def test_reload_w_bound_client(self): ) as final_attributes: job.reload() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None ) @@ -3668,7 +3649,7 @@ def test_reload_w_alternate_client(self): ) as final_attributes: job.reload(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5012,7 +4993,7 @@ def test__begin_w_timeout(self): ) as final_attributes: job._begin(timeout=7.5) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="POST", @@ -5049,7 +5030,7 @@ def test_begin_w_bound_client(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) self.assertIsNone(job.default_dataset) self.assertEqual(job.udf_resources, []) @@ -5133,7 +5114,7 @@ def test_begin_w_alternate_client(self): ) as final_attributes: job._begin(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5179,7 +5160,7 @@ def test_begin_w_udf(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) self.assertEqual(job.udf_resources, udf_resources) conn.api_request.assert_called_once_with( @@ -5233,7 +5214,7 @@ def test_begin_w_named_query_parameter(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) self.assertEqual(job.query_parameters, query_parameters) conn.api_request.assert_called_once_with( @@ -5281,7 +5262,7 @@ def test_begin_w_positional_query_parameter(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) self.assertEqual(job.query_parameters, query_parameters) conn.api_request.assert_called_once_with( @@ -5361,7 +5342,7 @@ def test_begin_w_table_defs(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="POST", @@ -5404,7 +5385,7 @@ def test_dry_run_query(self): ) as final_attributes: job._begin() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) self.assertEqual(job.udf_resources, []) conn.api_request.assert_called_once_with( method="POST", @@ -5430,7 +5411,7 @@ def test_exists_miss_w_bound_client(self): ) as final_attributes: self.assertFalse(job.exists()) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None @@ -5448,7 +5429,7 @@ def test_exists_hit_w_alternate_client(self): ) as final_attributes: self.assertTrue(job.exists(client=client2)) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5475,7 +5456,7 @@ def test_reload_w_bound_client(self): ) as final_attributes: job.reload() - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) self.assertNotEqual(job.destination, table_ref) @@ -5505,7 +5486,7 @@ def test_reload_w_alternate_client(self): ) as final_attributes: job.reload(client=client2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client2, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5533,7 +5514,7 @@ def test_reload_w_timeout(self): ) as final_attributes: job.reload(timeout=4.2) - final_attributes.assert_called_with({"path": PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, client, job) self.assertNotEqual(job.destination, table_ref)