From b13ad4749f7ae60017e3e030e31ae6846d417f86 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Sat, 1 Aug 2020 17:41:43 -0400 Subject: [PATCH 01/41] testing first trace export --- google/cloud/bigquery/client.py | 10 ++- .../cloud/bigquery/opentelemetry_tracing.py | 74 +++++++++++++++++++ 2 files changed, 80 insertions(+), 4 deletions(-) create mode 100644 google/cloud/bigquery/opentelemetry_tracing.py diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 651f0263e..c2c989ad2 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -33,6 +33,7 @@ import tempfile import uuid import warnings +from opentelemetry_tracing import SpanCreator try: import pyarrow @@ -2523,10 +2524,11 @@ def query( google.cloud.bigquery.job.QueryJobConfig, ) job_config = copy.deepcopy(self._default_query_job_config) - - job_ref = job._JobReference(job_id, project=project, location=location) - query_job = job.QueryJob(job_ref, query, client=self, job_config=job_config) - query_job._begin(retry=retry, timeout=timeout) + x = SpanCreator + with x.create(): + job_ref = job._JobReference(job_id, project=project, location=location) + query_job = job.QueryJob(job_ref, query, client=self, job_config=job_config) + query_job._begin(retry=retry, timeout=timeout) return query_job diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py new file mode 100644 index 000000000..62f1fb84c --- /dev/null +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -0,0 +1,74 @@ +# Copyright 2020 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from contextlib import contextmanager + +from google.api_core.exceptions import GoogleAPICallError +import pdb + +Logger = logging.getLogger(__name__) + +try: + from opentelemetry import trace + from opentelemetry import propagators + from opentelemetry.trace import SpanContext + from opentelemetry.trace import get_current_span + from opentelemetry.trace import set_span_in_context + from opentelemetry.trace.status import Status + from opentelemetry.instrumentation.utils import http_status_to_canonical_code + + HAS_OPENTELEMETRY = True + +except ImportError: + Logger.info( + 'This service instrumented using opentelemetry.' + 'Opentelemetry could not be imported please' + 'add opentelemetry-api and opentelemetry-instrumentation' + 'packages in order to get Big Query Tracing data' + ) + + HAS_OPENTELEMETRY = False + + +class SpanCreator: + def __init__(self): + self.opentelemetry_enbled = HAS_OPENTELEMETRY + + @contextmanager + def create(self, name, attributes=None, parent_context=None): + + print('yerr') + + if not self.opentelemetry_enbled: + yield None + return + + tracer = trace.get_tracer(__name__) + + default_attributes = { + 'db.type': 'BigQuery', + } + + if attributes: + default_attributes.update(attributes) + + # yield new span value + with tracer.start_as_current_span(name=name, attributes=attributes, parent=parent_context) as span: + try: + yield span + except GoogleAPICallError as error: + if error.code is not None: + span.set_status(Status(http_status_to_canonical_code(error.code))) + raise From cacea7cc500aaf796329ef57eeb2999dd275e78f Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Mon, 3 Aug 2020 01:07:35 -0400 Subject: [PATCH 02/41] instrumention client.py --- google/cloud/bigquery/client.py | 690 ++++++++++++++++++++------------ 1 file changed, 423 insertions(+), 267 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index c2c989ad2..222114aa9 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -80,7 +80,6 @@ from google.cloud.bigquery.table import TableReference from google.cloud.bigquery.table import RowIterator - _DEFAULT_CHUNKSIZE = 1048576 # 1024 * 1024 B = 1 MB _MAX_MULTIPART_SIZE = 5 * 1024 * 1024 _DEFAULT_NUM_RETRIES = 6 @@ -168,14 +167,14 @@ class Client(ClientWithProject): """The scopes required for authenticating as a BigQuery consumer.""" def __init__( - self, - project=None, - credentials=None, - _http=None, - location=None, - default_query_job_config=None, - client_info=None, - client_options=None, + self, + project=None, + credentials=None, + _http=None, + location=None, + default_query_job_config=None, + client_info=None, + client_options=None, ): super(Client, self).__init__( project=project, credentials=credentials, _http=_http @@ -213,7 +212,7 @@ def close(self): self._http.close() def get_service_account_email( - self, project=None, retry=DEFAULT_RETRY, timeout=None + self, project=None, retry=DEFAULT_RETRY, timeout=None ): """Get the email address of the project's BigQuery service account @@ -244,12 +243,14 @@ def get_service_account_email( if project is None: project = self.project path = "/projects/%s/serviceAccount" % (project,) - - api_response = self._call_api(retry, method="GET", path=path, timeout=timeout) + span_creator = SpanCreator() + attributes = {'path': path} + with span_creator.create('BigQuery.getServiceAccountEmail', attributes): + api_response = self._call_api(retry, method="GET", path=path, timeout=timeout) return api_response["email"] def list_projects( - self, max_results=None, page_token=None, retry=DEFAULT_RETRY, timeout=None + self, max_results=None, page_token=None, retry=DEFAULT_RETRY, timeout=None ): """List projects for the project associated with this client. @@ -279,10 +280,15 @@ def list_projects( Iterator of :class:`~google.cloud.bigquery.client.Project` accessible to the current client. """ + path = '/projects' + span_creator = SpanCreator() + attributes = {'path': path} + with span_creator.create('BigQuery.list_projects', attributes): + api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, - api_request=functools.partial(self._call_api, retry, timeout=timeout), - path="/projects", + api_request=api_request, + path=path, item_to_value=_item_to_project, items_key="projects", page_token=page_token, @@ -290,14 +296,14 @@ def list_projects( ) def list_datasets( - self, - project=None, - include_all=False, - filter=None, - max_results=None, - page_token=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + project=None, + include_all=False, + filter=None, + max_results=None, + page_token=None, + retry=DEFAULT_RETRY, + timeout=None, ): """List datasets for the project associated with this client. @@ -343,9 +349,15 @@ def list_datasets( # and converting it into a string here. extra_params["filter"] = filter path = "/projects/%s/datasets" % (project,) + span_creator = SpanCreator() + attributes = extra_params + attributes['path'] = path + attributes['page_token'] = page_token + with span_creator.create('BigQuery.list_datasets', attributes): + api_request = functools.partial(self._call_api, retry, timeout=timeout) 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", @@ -414,7 +426,7 @@ def _create_bqstorage_client(self): return bigquery_storage_v1.BigQueryReadClient(credentials=self._credentials) def create_dataset( - self, dataset, exists_ok=False, retry=DEFAULT_RETRY, timeout=None + self, dataset, exists_ok=False, retry=DEFAULT_RETRY, timeout=None ): """API call: create the dataset via a POST request. @@ -469,9 +481,13 @@ def create_dataset( data["location"] = self.location try: - api_response = self._call_api( - retry, method="POST", path=path, data=data, timeout=timeout - ) + span_creator = SpanCreator() + attributes = data + attributes['path'] = path + with span_creator.create('BigQuery.createDataset', attributes): + api_response = self._call_api( + retry, 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: @@ -479,7 +495,7 @@ def create_dataset( return self.get_dataset(dataset.reference, retry=retry) def create_routine( - self, routine, exists_ok=False, retry=DEFAULT_RETRY, timeout=None + self, routine, exists_ok=False, retry=DEFAULT_RETRY, timeout=None ): """[Beta] Create a routine via a POST request. @@ -513,9 +529,13 @@ def create_routine( ) resource = routine.to_api_repr() try: - api_response = self._call_api( - retry, method="POST", path=path, data=resource, timeout=timeout - ) + span_creator = SpanCreator() + attributes = {'path': path, + 'exists_ok': exists_ok} + with span_creator.create('BigQuery.createRoutine', attributes): + api_response = self._call_api( + retry, 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: @@ -556,13 +576,17 @@ def create_table(self, table, exists_ok=False, retry=DEFAULT_RETRY, timeout=None If the table already exists. """ table = _table_arg_to_table(table, default_project=self.project) - - path = "/projects/%s/datasets/%s/tables" % (table.project, table.dataset_id) + dataset_id = table.dataset_id + path = "/projects/%s/datasets/%s/tables" % (table.project, dataset_id) data = table.to_api_repr() try: - api_response = self._call_api( - retry, method="POST", path=path, data=data, timeout=timeout - ) + span_creator = SpanCreator() + attributes = {'path': path, + 'dataset_id': dataset_id} + with span_creator.create('BigQuery.createTable', attributes): + api_response = self._call_api( + retry, 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: @@ -597,18 +621,24 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None): google.cloud.bigquery.dataset.Dataset: A ``Dataset`` instance. """ + if isinstance(dataset_ref, str): + ref = dataset_ref dataset_ref = DatasetReference.from_string( dataset_ref, default_project=self.project ) - - api_response = self._call_api( - retry, method="GET", path=dataset_ref.path, timeout=timeout - ) + span_creator = SpanCreator() + path = dataset_ref.path + attributes = {'path': path, + 'dataset_ref': ref} + with span_creator.create('BigQuery.createTable', attributes): + api_response = self._call_api( + retry, method="GET", path=path, timeout=timeout + ) return Dataset.from_api_repr(api_response) def get_iam_policy( - self, table, requested_policy_version=1, retry=DEFAULT_RETRY, timeout=None, + self, table, requested_policy_version=1, retry=DEFAULT_RETRY, timeout=None, ): if not isinstance(table, (Table, TableReference)): raise TypeError("table must be a Table or TableReference") @@ -619,15 +649,18 @@ def get_iam_policy( body = {"options": {"requestedPolicyVersion": 1}} path = "{}:getIamPolicy".format(table.path) - - response = self._call_api( - retry, method="POST", path=path, data=body, timeout=timeout, - ) + span_creator = SpanCreator() + attributes = {'path': path, + 'body': body} + with span_creator.create('BigQuery.getIAMPolicy', attributes): + response = self._call_api( + retry, method="POST", path=path, data=body, timeout=timeout, + ) return Policy.from_api_repr(response) def set_iam_policy( - self, table, policy, updateMask=None, retry=DEFAULT_RETRY, timeout=None, + self, table, policy, updateMask=None, retry=DEFAULT_RETRY, timeout=None, ): if not isinstance(table, (Table, TableReference)): raise TypeError("table must be a Table or TableReference") @@ -641,15 +674,18 @@ def set_iam_policy( body["updateMask"] = updateMask path = "{}:setIamPolicy".format(table.path) - - response = self._call_api( - retry, method="POST", path=path, data=body, timeout=timeout, - ) + span_creator = SpanCreator() + attributes = {'path': path, + 'body': body} + with span_creator.create('BigQuery.setIAMPolicy', attributes): + response = self._call_api( + retry, method="POST", path=path, data=body, timeout=timeout, + ) return Policy.from_api_repr(response) def test_iam_permissions( - self, table, permissions, retry=DEFAULT_RETRY, timeout=None, + self, table, permissions, retry=DEFAULT_RETRY, timeout=None, ): if not isinstance(table, (Table, TableReference)): raise TypeError("table must be a Table or TableReference") @@ -657,10 +693,13 @@ def test_iam_permissions( body = {"permissions": permissions} path = "{}:testIamPermissions".format(table.path) - - response = self._call_api( - retry, method="POST", path=path, data=body, timeout=timeout, - ) + span_creator = SpanCreator() + attributes = {'path': path, + 'body': body} + with span_creator.create('BigQuery.testIAMPolicy', attributes): + response = self._call_api( + retry, method="POST", path=path, data=body, timeout=timeout, + ) return response @@ -686,13 +725,18 @@ def get_model(self, model_ref, retry=DEFAULT_RETRY, timeout=None): google.cloud.bigquery.model.Model: A ``Model`` instance. """ if isinstance(model_ref, str): + ref = model_ref model_ref = ModelReference.from_string( model_ref, default_project=self.project ) - - api_response = self._call_api( - retry, method="GET", path=model_ref.path, timeout=timeout - ) + span_creator = SpanCreator() + path = model_ref.path + attributes = {'path': path, + 'model_ref': ref} + with span_creator.create('BigQuery.getModel', attributes): + api_response = self._call_api( + retry, method="GET", path=path, timeout=timeout + ) return Model.from_api_repr(api_response) def get_routine(self, routine_ref, retry=DEFAULT_RETRY, timeout=None): @@ -718,14 +762,20 @@ def get_routine(self, routine_ref, retry=DEFAULT_RETRY, timeout=None): google.cloud.bigquery.routine.Routine: A ``Routine`` instance. """ + if isinstance(routine_ref, str): + ref = routine_ref routine_ref = RoutineReference.from_string( routine_ref, default_project=self.project ) - - api_response = self._call_api( - retry, method="GET", path=routine_ref.path, timeout=timeout - ) + path = routine_ref.path + span_creator = SpanCreator() + attributes = {'path': path, + 'routine_ref': ref} + with span_creator.create('BigQuery.getRoutine', attributes): + api_response = self._call_api( + retry, method="GET", path=path, timeout=timeout + ) return Routine.from_api_repr(api_response) def get_table(self, table, retry=DEFAULT_RETRY, timeout=None): @@ -752,9 +802,13 @@ def get_table(self, table, retry=DEFAULT_RETRY, timeout=None): A ``Table`` instance. """ table_ref = _table_arg_to_table_ref(table, default_project=self.project) - api_response = self._call_api( - retry, method="GET", path=table_ref.path, timeout=timeout - ) + path = table_ref.path + span_creator = SpanCreator() + attributes = {'path': path} + with span_creator.create('BigQuery.getTable', attributes): + api_response = self._call_api( + retry, method="GET", path=path, timeout=timeout + ) return Table.from_api_repr(api_response) def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): @@ -791,14 +845,20 @@ def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): headers = {"If-Match": dataset.etag} else: headers = None - api_response = self._call_api( - retry, - method="PATCH", - path=dataset.path, - data=partial, - headers=headers, - timeout=timeout, - ) + path = dataset.path + span_creator = SpanCreator() + attributes = headers + attributes['path'] = path + attributes['fields'] = fields + with span_creator.create('BigQuery.updateDataset', attributes): + api_response = self._call_api( + retry, + method="PATCH", + path=dataset.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): @@ -834,14 +894,20 @@ def update_model(self, model, fields, retry=DEFAULT_RETRY, timeout=None): headers = {"If-Match": model.etag} else: headers = None - api_response = self._call_api( - retry, - method="PATCH", - path=model.path, - data=partial, - headers=headers, - timeout=timeout, - ) + path = model.path + span_creator = SpanCreator() + attributes = headers + attributes['path'] = path + attributes['fields'] = fields + with span_creator.create('BigQuery.updateModel', attributes): + api_response = self._call_api( + retry, + 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): @@ -888,14 +954,20 @@ def update_routine(self, routine, fields, retry=DEFAULT_RETRY, timeout=None): # TODO: remove when routines update supports partial requests. partial["routineReference"] = routine.reference.to_api_repr() - api_response = self._call_api( - retry, - method="PUT", - path=routine.path, - data=partial, - headers=headers, - timeout=timeout, - ) + path = routine.path + span_creator = SpanCreator() + attributes = headers + attributes['path'] = path + attributes['fields'] = fields + with span_creator.create('BigQuery.updateRountine', attributes): + api_response = self._call_api( + retry, + 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): @@ -931,23 +1003,30 @@ def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): headers = {"If-Match": table.etag} else: headers = None - api_response = self._call_api( - retry, - method="PATCH", - path=table.path, - data=partial, - headers=headers, - timeout=timeout, - ) + + path = table.path + span_creator = SpanCreator() + attributes = headers + attributes['path'] = path + attributes['fields'] = fields + with span_creator.create('BigQuery.updateTable', attributes): + api_response = self._call_api( + retry, + method="PATCH", + path=table.path, + data=partial, + headers=headers, + timeout=timeout, + ) return Table.from_api_repr(api_response) def list_models( - self, - dataset, - max_results=None, - page_token=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + dataset, + max_results=None, + page_token=None, + retry=DEFAULT_RETRY, + timeout=None, ): """[Beta] List models in the dataset. @@ -994,9 +1073,15 @@ def list_models( raise TypeError("dataset must be a Dataset, DatasetReference, or string") path = "%s/models" % dataset.path + span_creator = SpanCreator() + attributes = {'path': path, + 'page_token': page_token, + 'max_results': max_results} + with span_creator.create('BigQuery.listModels', attributes): + api_request = functools.partial(self._call_api, retry, timeout=timeout) 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", @@ -1007,12 +1092,12 @@ def list_models( return result def list_routines( - self, - dataset, - max_results=None, - page_token=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + dataset, + max_results=None, + page_token=None, + retry=DEFAULT_RETRY, + timeout=None, ): """[Beta] List routines in the dataset. @@ -1059,9 +1144,15 @@ def list_routines( raise TypeError("dataset must be a Dataset, DatasetReference, or string") path = "{}/routines".format(dataset.path) + span_creator = SpanCreator() + attributes = {'path': path, + 'page_token': page_token, + 'max_results': max_results} + with span_creator.create('BigQuery.listRoutines', attributes): + api_requests = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( client=self, - api_request=functools.partial(self._call_api, retry, timeout=timeout), + api_request=api_requests, path=path, item_to_value=_item_to_routine, items_key="routines", @@ -1072,12 +1163,12 @@ def list_routines( return result def list_tables( - self, - dataset, - max_results=None, - page_token=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + dataset, + max_results=None, + page_token=None, + retry=DEFAULT_RETRY, + timeout=None, ): """List tables in the dataset. @@ -1124,9 +1215,16 @@ def list_tables( raise TypeError("dataset must be a Dataset, DatasetReference, or string") path = "%s/tables" % dataset.path + span_creator = SpanCreator() + attributes = {'path': path, + 'page_token': page_token, + 'max_results': max_results} + with span_creator.create('BigQuery.listTables', attributes): + api_requests = functools.partial(self._call_api, retry, timeout=timeout) + result = page_iterator.HTTPIterator( client=self, - api_request=functools.partial(self._call_api, retry, timeout=timeout), + api_request=api_requests, path=path, item_to_value=_item_to_table, items_key="tables", @@ -1137,12 +1235,12 @@ def list_tables( return result def delete_dataset( - self, - dataset, - delete_contents=False, - retry=DEFAULT_RETRY, - timeout=None, - not_found_ok=False, + self, + dataset, + delete_contents=False, + retry=DEFAULT_RETRY, + timeout=None, + not_found_ok=False, ): """Delete a dataset. @@ -1181,23 +1279,28 @@ def delete_dataset( raise TypeError("dataset must be a Dataset or a DatasetReference") params = {} + path = dataset.path if delete_contents: params["deleteContents"] = "true" try: - self._call_api( - retry, - method="DELETE", - path=dataset.path, - query_params=params, - timeout=timeout, - ) + span_creator = SpanCreator() + attributes = params + attributes["path"] = path + with span_creator.create('BigQuery.deleteDataset', attributes): + self._call_api( + retry, + method="DELETE", + path=path, + query_params=params, + timeout=timeout, + ) except google.api_core.exceptions.NotFound: if not not_found_ok: raise def delete_model( - self, model, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False + self, model, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False ): """[Beta] Delete a model @@ -1224,19 +1327,26 @@ def delete_model( when deleting the model. """ if isinstance(model, str): + ref = model model = ModelReference.from_string(model, default_project=self.project) if not isinstance(model, (Model, ModelReference)): raise TypeError("model must be a Model or a ModelReference") + path = model.path try: - self._call_api(retry, method="DELETE", path=model.path, timeout=timeout) + span_creator = SpanCreator() + attributes = {"model": ref, + "not_found_okay": not_found_ok, + "path": path} + with span_creator.create('BigQuery.deleteModel', attributes): + self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: raise def delete_routine( - self, routine, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False + self, routine, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False ): """[Beta] Delete a routine. @@ -1263,21 +1373,27 @@ def delete_routine( when deleting the routine. """ if isinstance(routine, str): + ref = routine routine = RoutineReference.from_string( routine, default_project=self.project ) + path = routine.path if not isinstance(routine, (Routine, RoutineReference)): raise TypeError("routine must be a Routine or a RoutineReference") try: - self._call_api(retry, method="DELETE", path=routine.path, timeout=timeout) + span_creator = SpanCreator() + attributes = {"path": path, + "routine": ref} + with span_creator.create('BigQuery.deleteRoutine', attributes): + self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: raise def delete_table( - self, table, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False + self, table, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False ): """Delete a table @@ -1308,13 +1424,17 @@ def delete_table( raise TypeError("Unable to get TableReference for table '{}'".format(table)) try: - self._call_api(retry, method="DELETE", path=table.path, timeout=timeout) + path = table.path + span_creator = SpanCreator() + attributes = {"path": path} + with span_creator.create('BigQuery.deleteTable', attributes): + self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: raise def _get_query_results( - self, job_id, retry, project=None, timeout_ms=None, location=None, timeout=None + self, job_id, retry, project=None, timeout_ms=None, location=None, timeout=None ): """Get the query results object for a query job. @@ -1356,9 +1476,15 @@ def _get_query_results( # This call is typically made in a polling loop that checks whether the # job is complete (from QueryJob.done(), called ultimately from # QueryJob.result()). So we don't need to poll here. - resource = self._call_api( - retry, method="GET", path=path, query_params=extra_params, timeout=timeout - ) + span_creator = SpanCreator() + attributes = {"path": path, + "job_id": job_id, + "location": location} + + with span_creator.create('BigQuery.getQueryResults', attributes): + resource = self._call_api( + retry, method="GET", path=path, query_params=extra_params, timeout=timeout + ) return _QueryResults.from_api_repr(resource) def job_from_resource(self, resource): @@ -1460,7 +1586,7 @@ def create_job(self, job_config, retry=DEFAULT_RETRY): raise TypeError("Invalid job configuration received.") def get_job( - self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None + self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None ): """Fetch a job for the project associated with this client. @@ -1501,15 +1627,20 @@ def get_job( extra_params["location"] = location path = "/projects/{}/jobs/{}".format(project, job_id) - - resource = self._call_api( - retry, method="GET", path=path, query_params=extra_params, timeout=timeout - ) + span_creator = SpanCreator() + attributes = extra_params + attributes["path"] = path + attributes["job_id"] = job_id + attributes["location"] = location + with span_creator.create('BigQuery.getJob', attributes): + resource = self._call_api( + retry, method="GET", path=path, query_params=extra_params, timeout=timeout + ) return self.job_from_resource(resource) def cancel_job( - self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None + self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None ): """Attempt to cancel a job from a job ID. @@ -1550,25 +1681,30 @@ def cancel_job( extra_params["location"] = location path = "/projects/{}/jobs/{}/cancel".format(project, job_id) - - resource = self._call_api( - retry, method="POST", path=path, query_params=extra_params, timeout=timeout - ) + span_creator = SpanCreator() + attributes = extra_params + attributes["path"] = path + attributes["job_id"] = job_id + attributes["location"] = location + with span_creator.create('BigQuery.CancelJob', attributes): + resource = self._call_api( + retry, method="POST", path=path, query_params=extra_params, timeout=timeout + ) return self.job_from_resource(resource["job"]) def list_jobs( - self, - project=None, - parent_job=None, - max_results=None, - page_token=None, - all_users=None, - state_filter=None, - retry=DEFAULT_RETRY, - timeout=None, - min_creation_time=None, - max_creation_time=None, + self, + project=None, + parent_job=None, + max_results=None, + page_token=None, + all_users=None, + state_filter=None, + retry=DEFAULT_RETRY, + timeout=None, + min_creation_time=None, + max_creation_time=None, ): """List jobs for the project associated with this client. @@ -1643,9 +1779,16 @@ def list_jobs( project = self.project path = "/projects/%s/jobs" % (project,) + span_creator = SpanCreator() + attributes = extra_params + attributes["path"] = path + attributes["all_users"] = all_users + attributes["state_filter"] = state_filter + with span_creator.create('BigQuery.listJobs', attributes): + api_request = functools.partial(self._call_api, retry, timeout=timeout) 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", @@ -1655,16 +1798,16 @@ def list_jobs( ) def load_table_from_uri( - self, - source_uris, - destination, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + source_uris, + destination, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, + retry=DEFAULT_RETRY, + timeout=None, ): """Starts a job for loading data into a table from CloudStorage. @@ -1737,17 +1880,17 @@ def load_table_from_uri( return load_job def load_table_from_file( - self, - file_obj, - destination, - rewind=False, - size=None, - num_retries=_DEFAULT_NUM_RETRIES, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, + self, + file_obj, + destination, + rewind=False, + size=None, + num_retries=_DEFAULT_NUM_RETRIES, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, ): """Upload the contents of this table from a file-like object. @@ -1837,16 +1980,16 @@ def load_table_from_file( return self.job_from_resource(response.json()) def load_table_from_dataframe( - self, - dataframe, - destination, - num_retries=_DEFAULT_NUM_RETRIES, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, - parquet_compression="snappy", + self, + dataframe, + destination, + num_retries=_DEFAULT_NUM_RETRIES, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, + parquet_compression="snappy", ): """Upload the contents of a table from a pandas DataFrame. @@ -1945,8 +2088,8 @@ def load_table_from_dataframe( # schema, and check if dataframe schema is compatible with it - except # for WRITE_TRUNCATE jobs, the existing schema does not matter then. if ( - not job_config.schema - and job_config.write_disposition != job.WriteDisposition.WRITE_TRUNCATE + not job_config.schema + and job_config.write_disposition != job.WriteDisposition.WRITE_TRUNCATE ): try: table = self.get_table(destination) @@ -2028,15 +2171,15 @@ def load_table_from_dataframe( os.remove(tmppath) def load_table_from_json( - self, - json_rows, - destination, - num_retries=_DEFAULT_NUM_RETRIES, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, + self, + json_rows, + destination, + num_retries=_DEFAULT_NUM_RETRIES, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, ): """Upload the contents of a table from a JSON string or dict. @@ -2243,16 +2386,16 @@ def _do_multipart_upload(self, stream, metadata, size, num_retries): return response def copy_table( - self, - sources, - destination, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + sources, + destination, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, + retry=DEFAULT_RETRY, + timeout=None, ): """Copy one or more tables to another table. @@ -2346,17 +2489,17 @@ def copy_table( return copy_job def extract_table( - self, - source, - destination_uris, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, - retry=DEFAULT_RETRY, - timeout=None, - source_type="Table", + self, + source, + destination_uris, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, + retry=DEFAULT_RETRY, + timeout=None, + source_type="Table", ): """Start a job to extract a table into Cloud Storage files. @@ -2445,15 +2588,15 @@ def extract_table( return extract_job def query( - self, - query, - job_config=None, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + query, + job_config=None, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + retry=DEFAULT_RETRY, + timeout=None, ): """Run a SQL query. @@ -2595,7 +2738,7 @@ def insert_rows(self, table, rows, selected_fields=None, **kwargs): return self.insert_rows_json(table, json_rows, **kwargs) def insert_rows_from_dataframe( - self, table, dataframe, selected_fields=None, chunk_size=500, **kwargs + self, table, dataframe, selected_fields=None, chunk_size=500, **kwargs ): """Insert rows into a table from a dataframe via the streaming API. @@ -2643,15 +2786,15 @@ def insert_rows_from_dataframe( return insert_results def insert_rows_json( - self, - table, - json_rows, - row_ids=None, - skip_invalid_rows=None, - ignore_unknown_values=None, - template_suffix=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + table, + json_rows, + row_ids=None, + skip_invalid_rows=None, + ignore_unknown_values=None, + template_suffix=None, + retry=DEFAULT_RETRY, + timeout=None, ): """Insert rows into a table without applying local type conversions. @@ -2702,7 +2845,7 @@ def insert_rows_json( TypeError: if `json_rows` is not a `Sequence`. """ if not isinstance( - json_rows, (collections_abc.Sequence, collections_abc.Iterator) + json_rows, (collections_abc.Sequence, collections_abc.Iterator) ): raise TypeError("json_rows argument should be a sequence of dicts") # Convert table to just a reference because unlike insert_rows, @@ -2729,14 +2872,19 @@ def insert_rows_json( if template_suffix is not None: data["templateSuffix"] = template_suffix + path = "%s/insertAll" % table.path # We can always retry, because every row has an insert ID. - response = self._call_api( - retry, - method="POST", - path="%s/insertAll" % table.path, - data=data, - timeout=timeout, - ) + span_creator = SpanCreator() + attributes = data + attributes["path"] = path + with span_creator.create('BigQuery.insertRowsJson', attributes): + response = self._call_api( + retry, + method="POST", + path=path, + data=data, + timeout=timeout, + ) errors = [] for error in response.get("insertErrors", ()): @@ -2785,15 +2933,15 @@ def list_partitions(self, table, retry=DEFAULT_RETRY, timeout=None): ] def list_rows( - self, - table, - selected_fields=None, - max_results=None, - page_token=None, - start_index=None, - page_size=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + table, + selected_fields=None, + max_results=None, + page_token=None, + start_index=None, + page_size=None, + retry=DEFAULT_RETRY, + timeout=None, ): """List the rows of the table. @@ -2874,11 +3022,19 @@ def list_rows( params["selectedFields"] = ",".join(field.name for field in selected_fields) if start_index is not None: params["startIndex"] = start_index - + span_creator = SpanCreator() + path = "%s/data" % (table.path,) + attributes = params + attributes["path"] = path + attributes["page_size"] = page_size + attributes["page_token"] = page_token + attributes["start_index"] = start_index + with span_creator.create('BigQuery.listRows', attributes): + api_request = functools.partial(self._call_api, retry, timeout=timeout) row_iterator = RowIterator( client=self, - api_request=functools.partial(self._call_api, retry, timeout=timeout), - path="%s/data" % (table.path,), + api_request=api_request, + path=path, schema=schema, page_token=page_token, max_results=max_results, From 63e9581f90ef33f946cb2938643379ad5770d0ff Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Mon, 3 Aug 2020 23:29:22 -0400 Subject: [PATCH 03/41] instrumenting job.py and adding documentation --- README.rst | 48 ++++++++++++++++++++++++++ google/cloud/bigquery/job.py | 67 ++++++++++++++++++++++-------------- 2 files changed, 90 insertions(+), 25 deletions(-) diff --git a/README.rst b/README.rst index 8f73576d6..7141b8456 100644 --- a/README.rst +++ b/README.rst @@ -102,3 +102,51 @@ Perform a query for row in rows: print(row.name) + +Instrumenting With OpenTelemetry +-------------------------------- + +This application uses OpenTelmetry to output tracing data from +api calls to BigQuery. To enable OpenTelemetry tracing in +the BigQuery client the following PyPi packages need to be installed: + +.. code-block:: console + + pip install opentelemetry-api + +.. code-block:: console + + pip install opentelemetry-sdk + +.. code-block:: console + + pip install opentelemetry-instrumentaion + +.. code-block:: console + + pip install opentelemetry-exporter-cloud_trace + +After installation, OpenTelemetry can now be used in the BigQuery +client and in BigQuery jobs. First however an exporter must be +specified for where the trace data will be outputted to. An example of this +can be found here: + +.. code-block:: python + + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor + from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter + trace.set_tracer_provider(TracerProvider()) + trace.get_tracer_provider().add_span_processor( + SimpleExportSpanProcessor(CloudTraceSpanExporter()) + ) + +In this example all tracing data will be published to the Google +Cloud trace console. For more information on OpenTelemetry, please consult the `OpenTelemetry documentation`_. + +.. _OpenTelemetry documentation: https://opentelemetry-python.readthedocs.io + + + + diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index 753307b2a..712aca50f 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -50,6 +50,8 @@ from google.cloud.bigquery.table import TableReference from google.cloud.bigquery.table import Table from google.cloud.bigquery.table import TimePartitioning +from opentelemetry_tracing import SpanCreator + _DONE_STATE = "DONE" _STOPPED_REASON = "stopped" @@ -634,9 +636,12 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None): # jobs.insert is idempotent because we ensure that every new # job has an ID. - api_response = client._call_api( - retry, method="POST", path=path, data=self.to_api_repr(), timeout=timeout - ) + attributes = {'path': path} + span_creator = SpanCreator() + with span_creator.create('BigQuery.job.begin', attributes): + api_response = client._call_api( + retry, 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): @@ -665,13 +670,17 @@ def exists(self, client=None, retry=DEFAULT_RETRY, timeout=None): extra_params["location"] = self.location try: - client._call_api( - retry, - method="GET", - path=self.path, - query_params=extra_params, - timeout=timeout, - ) + span_creator = SpanCreator() + attributes = extra_params + attributes['path'] = self.path + with span_creator.create('BigQuery.job.exists', attributes): + client._call_api( + retry, + method="GET", + path=self.path, + query_params=extra_params, + timeout=timeout, + ) except NotFound: return False else: @@ -698,14 +707,17 @@ def reload(self, client=None, retry=DEFAULT_RETRY, timeout=None): extra_params = {} if self.location: extra_params["location"] = self.location - - api_response = client._call_api( - retry, - method="GET", - path=self.path, - query_params=extra_params, - timeout=timeout, - ) + span_creator = SpanCreator() + attributes = extra_params + attributes['path'] = self.path + with span_creator.create('BigQuery.job.reload', attributes): + api_response = client._call_api( + retry, + 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): @@ -732,13 +744,18 @@ def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None): if self.location: extra_params["location"] = self.location - api_response = client._call_api( - retry, - method="POST", - path="{}/cancel".format(self.path), - query_params=extra_params, - timeout=timeout, - ) + path = "{}/cancel".format(self.path) + span_creator = SpanCreator() + attributes = extra_params + attributes['path'] = path + with span_creator.create('BigQuery.job.cancel', attributes): + api_response = client._call_api( + retry, + 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. From 719ed83d3c0d4f4f03a060fb5c2c788cda72ff84 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Tue, 4 Aug 2020 19:28:30 -0400 Subject: [PATCH 04/41] reconfiguring imports --- google/cloud/bigquery/client.py | 11 +++++------ google/cloud/bigquery/job.py | 2 +- google/cloud/bigquery/opentelemetry_tracing.py | 3 --- 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 222114aa9..f78081736 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -33,7 +33,7 @@ import tempfile import uuid import warnings -from opentelemetry_tracing import SpanCreator +from google.cloud.bigquery.opentelemetry_tracing import SpanCreator try: import pyarrow @@ -2667,11 +2667,10 @@ def query( google.cloud.bigquery.job.QueryJobConfig, ) job_config = copy.deepcopy(self._default_query_job_config) - x = SpanCreator - with x.create(): - job_ref = job._JobReference(job_id, project=project, location=location) - query_job = job.QueryJob(job_ref, query, client=self, job_config=job_config) - query_job._begin(retry=retry, timeout=timeout) + + job_ref = job._JobReference(job_id, project=project, location=location) + query_job = job.QueryJob(job_ref, query, client=self, job_config=job_config) + query_job._begin(retry=retry, timeout=timeout) return query_job diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index 712aca50f..0eec50506 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -50,7 +50,7 @@ from google.cloud.bigquery.table import TableReference from google.cloud.bigquery.table import Table from google.cloud.bigquery.table import TimePartitioning -from opentelemetry_tracing import SpanCreator +from google.cloud.bigquery.opentelemetry_tracing import SpanCreator _DONE_STATE = "DONE" diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 62f1fb84c..334910831 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -48,9 +48,6 @@ def __init__(self): @contextmanager def create(self, name, attributes=None, parent_context=None): - - print('yerr') - if not self.opentelemetry_enbled: yield None return From 72f22b6aeaf6a1883a29df686260baa89451fc58 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Wed, 5 Aug 2020 10:21:24 -0400 Subject: [PATCH 05/41] quick cleanup of unused variable --- google/cloud/bigquery/opentelemetry_tracing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 334910831..280e61d1f 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -47,7 +47,7 @@ def __init__(self): self.opentelemetry_enbled = HAS_OPENTELEMETRY @contextmanager - def create(self, name, attributes=None, parent_context=None): + def create(self, name, attributes=None): if not self.opentelemetry_enbled: yield None return @@ -62,7 +62,7 @@ def create(self, name, attributes=None, parent_context=None): default_attributes.update(attributes) # yield new span value - with tracer.start_as_current_span(name=name, attributes=attributes, parent=parent_context) as span: + with tracer.start_as_current_span(name=name, attributes=default_attributes) as span: try: yield span except GoogleAPICallError as error: From ff8171ffa8e0ea025bb7f5ec836279c40f3e2060 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Fri, 7 Aug 2020 14:45:29 -0400 Subject: [PATCH 06/41] adding more attributes in module and limiting complexity of instrumentation --- google/cloud/bigquery/client.py | 164 +++++++++--------- google/cloud/bigquery/job.py | 14 +- .../cloud/bigquery/opentelemetry_tracing.py | 43 ++++- 3 files changed, 123 insertions(+), 98 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index f78081736..34ee0ad7e 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -245,7 +245,7 @@ def get_service_account_email( path = "/projects/%s/serviceAccount" % (project,) span_creator = SpanCreator() attributes = {'path': path} - with span_creator.create('BigQuery.getServiceAccountEmail', attributes): + with span_creator.create(name='BigQuery.getServiceAccountEmail', attributes=attributes, client=self): api_response = self._call_api(retry, method="GET", path=path, timeout=timeout) return api_response["email"] @@ -283,7 +283,7 @@ def list_projects( path = '/projects' span_creator = SpanCreator() attributes = {'path': path} - with span_creator.create('BigQuery.list_projects', attributes): + with span_creator.create(name='BigQuery.list_projects', attributes=attributes, client=self): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, @@ -350,10 +350,11 @@ def list_datasets( extra_params["filter"] = filter path = "/projects/%s/datasets" % (project,) span_creator = SpanCreator() - attributes = extra_params - attributes['path'] = path - attributes['page_token'] = page_token - with span_creator.create('BigQuery.list_datasets', attributes): + attributes = {'path': path, + 'page_token': page_token + } + + with span_creator.create(name='BigQuery.list_datasets', attributes=attributes, client=self): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, @@ -482,9 +483,8 @@ def create_dataset( try: span_creator = SpanCreator() - attributes = data - attributes['path'] = path - with span_creator.create('BigQuery.createDataset', attributes): + attributes = {'path': path} + with span_creator.create(name='BigQuery.createDataset', attributes=attributes, client=self): api_response = self._call_api( retry, method="POST", path=path, data=data, timeout=timeout ) @@ -532,7 +532,7 @@ def create_routine( span_creator = SpanCreator() attributes = {'path': path, 'exists_ok': exists_ok} - with span_creator.create('BigQuery.createRoutine', attributes): + with span_creator.create(name='BigQuery.createRoutine', attributes=attributes, client=self): api_response = self._call_api( retry, method="POST", path=path, data=resource, timeout=timeout ) @@ -583,7 +583,7 @@ def create_table(self, table, exists_ok=False, retry=DEFAULT_RETRY, timeout=None span_creator = SpanCreator() attributes = {'path': path, 'dataset_id': dataset_id} - with span_creator.create('BigQuery.createTable', attributes): + with span_creator.create(name='BigQuery.createTable', attributes=attributes, client=self): api_response = self._call_api( retry, method="POST", path=path, data=data, timeout=timeout ) @@ -623,15 +623,13 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None): """ if isinstance(dataset_ref, str): - ref = dataset_ref dataset_ref = DatasetReference.from_string( dataset_ref, default_project=self.project ) span_creator = SpanCreator() path = dataset_ref.path - attributes = {'path': path, - 'dataset_ref': ref} - with span_creator.create('BigQuery.createTable', attributes): + attributes = {'path': path} + with span_creator.create(name='BigQuery.createTable', attributes=attributes, client=self): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout ) @@ -652,7 +650,7 @@ def get_iam_policy( span_creator = SpanCreator() attributes = {'path': path, 'body': body} - with span_creator.create('BigQuery.getIAMPolicy', attributes): + with span_creator.create(name='BigQuery.getIAMPolicy', attributes=attributes, client=self): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, ) @@ -677,7 +675,7 @@ def set_iam_policy( span_creator = SpanCreator() attributes = {'path': path, 'body': body} - with span_creator.create('BigQuery.setIAMPolicy', attributes): + with span_creator.create(name='BigQuery.setIAMPolicy', attributes=attributes, client=self): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, ) @@ -696,7 +694,7 @@ def test_iam_permissions( span_creator = SpanCreator() attributes = {'path': path, 'body': body} - with span_creator.create('BigQuery.testIAMPolicy', attributes): + with span_creator.create(name='BigQuery.testIAMPolicy', attributes=attributes, client=self): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, ) @@ -725,15 +723,13 @@ def get_model(self, model_ref, retry=DEFAULT_RETRY, timeout=None): google.cloud.bigquery.model.Model: A ``Model`` instance. """ if isinstance(model_ref, str): - ref = model_ref model_ref = ModelReference.from_string( model_ref, default_project=self.project ) span_creator = SpanCreator() path = model_ref.path - attributes = {'path': path, - 'model_ref': ref} - with span_creator.create('BigQuery.getModel', attributes): + attributes = {'path': path} + with span_creator.create(name='BigQuery.getModel', attributes=attributes, client=self): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout ) @@ -764,15 +760,13 @@ def get_routine(self, routine_ref, retry=DEFAULT_RETRY, timeout=None): """ if isinstance(routine_ref, str): - ref = routine_ref routine_ref = RoutineReference.from_string( routine_ref, default_project=self.project ) path = routine_ref.path span_creator = SpanCreator() - attributes = {'path': path, - 'routine_ref': ref} - with span_creator.create('BigQuery.getRoutine', attributes): + attributes = {'path': path} + with span_creator.create(name='BigQuery.getRoutine', attributes=attributes, client=self): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout ) @@ -805,7 +799,7 @@ def get_table(self, table, retry=DEFAULT_RETRY, timeout=None): path = table_ref.path span_creator = SpanCreator() attributes = {'path': path} - with span_creator.create('BigQuery.getTable', attributes): + with span_creator.create(name='BigQuery.getTable', attributes=attributes, client=self): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout ) @@ -847,14 +841,14 @@ def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): headers = None path = dataset.path span_creator = SpanCreator() - attributes = headers - attributes['path'] = path - attributes['fields'] = fields - with span_creator.create('BigQuery.updateDataset', attributes): + attributes = {'path': path, + 'fields': fields} + + with span_creator.create(name='BigQuery.updateDataset', attributes=attributes, client=self): api_response = self._call_api( retry, method="PATCH", - path=dataset.path, + path=path, data=partial, headers=headers, timeout=timeout, @@ -896,10 +890,10 @@ def update_model(self, model, fields, retry=DEFAULT_RETRY, timeout=None): headers = None path = model.path span_creator = SpanCreator() - attributes = headers - attributes['path'] = path - attributes['fields'] = fields - with span_creator.create('BigQuery.updateModel', attributes): + attributes = {'path': path, + 'fields': fields} + + with span_creator.create(name='BigQuery.updateModel', attributes=attributes, client=self): api_response = self._call_api( retry, method="PATCH", @@ -956,10 +950,10 @@ def update_routine(self, routine, fields, retry=DEFAULT_RETRY, timeout=None): path = routine.path span_creator = SpanCreator() - attributes = headers - attributes['path'] = path - attributes['fields'] = fields - with span_creator.create('BigQuery.updateRountine', attributes): + attributes = {'path': path, + 'fields': fields} + + with span_creator.create(name='BigQuery.updateRountine', attributes=attributes, client=self): api_response = self._call_api( retry, method="PUT", @@ -1006,14 +1000,14 @@ def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): path = table.path span_creator = SpanCreator() - attributes = headers - attributes['path'] = path - attributes['fields'] = fields - with span_creator.create('BigQuery.updateTable', attributes): + attributes = {'path': path, + 'fields': fields} + + with span_creator.create(name='BigQuery.updateTable', attributes=attributes, client=self): api_response = self._call_api( retry, method="PATCH", - path=table.path, + path=path, data=partial, headers=headers, timeout=timeout, @@ -1077,7 +1071,7 @@ def list_models( attributes = {'path': path, 'page_token': page_token, 'max_results': max_results} - with span_creator.create('BigQuery.listModels', attributes): + with span_creator.create(name='BigQuery.listModels', attributes=attributes, client=self): api_request = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( client=self, @@ -1148,7 +1142,7 @@ def list_routines( attributes = {'path': path, 'page_token': page_token, 'max_results': max_results} - with span_creator.create('BigQuery.listRoutines', attributes): + with span_creator.create(name='BigQuery.listRoutines', attributes=attributes, client=self): api_requests = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( client=self, @@ -1219,7 +1213,7 @@ def list_tables( attributes = {'path': path, 'page_token': page_token, 'max_results': max_results} - with span_creator.create('BigQuery.listTables', attributes): + with span_creator.create(name='BigQuery.listTables', attributes=attributes, client=self): api_requests = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( @@ -1285,9 +1279,8 @@ def delete_dataset( try: span_creator = SpanCreator() - attributes = params - attributes["path"] = path - with span_creator.create('BigQuery.deleteDataset', attributes): + attributes = {"path": path} + with span_creator.create(name='BigQuery.deleteDataset', attributes=attributes, client=self): self._call_api( retry, method="DELETE", @@ -1327,7 +1320,6 @@ def delete_model( when deleting the model. """ if isinstance(model, str): - ref = model model = ModelReference.from_string(model, default_project=self.project) if not isinstance(model, (Model, ModelReference)): @@ -1336,10 +1328,9 @@ def delete_model( path = model.path try: span_creator = SpanCreator() - attributes = {"model": ref, - "not_found_okay": not_found_ok, + attributes = {"not_found_okay": not_found_ok, "path": path} - with span_creator.create('BigQuery.deleteModel', attributes): + with span_creator.create(name='BigQuery.deleteModel', attributes=attributes, client=self): self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: @@ -1373,7 +1364,6 @@ def delete_routine( when deleting the routine. """ if isinstance(routine, str): - ref = routine routine = RoutineReference.from_string( routine, default_project=self.project ) @@ -1384,9 +1374,8 @@ def delete_routine( try: span_creator = SpanCreator() - attributes = {"path": path, - "routine": ref} - with span_creator.create('BigQuery.deleteRoutine', attributes): + attributes = {"path": path} + with span_creator.create(name='BigQuery.deleteRoutine', attributes=attributes, client=self): self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: @@ -1427,7 +1416,7 @@ def delete_table( path = table.path span_creator = SpanCreator() attributes = {"path": path} - with span_creator.create('BigQuery.deleteTable', attributes): + with span_creator.create(name='BigQuery.deleteTable', attributes=attributes, client=self): self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: @@ -1481,7 +1470,7 @@ def _get_query_results( "job_id": job_id, "location": location} - with span_creator.create('BigQuery.getQueryResults', attributes): + with span_creator.create(name='BigQuery.getQueryResults', attributes=attributes, client=self): resource = self._call_api( retry, method="GET", path=path, query_params=extra_params, timeout=timeout ) @@ -1628,11 +1617,13 @@ def get_job( path = "/projects/{}/jobs/{}".format(project, job_id) span_creator = SpanCreator() - attributes = extra_params - attributes["path"] = path - attributes["job_id"] = job_id - attributes["location"] = location - with span_creator.create('BigQuery.getJob', attributes): + + attributes = {"path": path, + "job_id": job_id, + "location": location + } + + with span_creator.create(name='BigQuery.getJob', attributes=attributes, client=self): resource = self._call_api( retry, method="GET", path=path, query_params=extra_params, timeout=timeout ) @@ -1682,11 +1673,13 @@ def cancel_job( path = "/projects/{}/jobs/{}/cancel".format(project, job_id) span_creator = SpanCreator() - attributes = extra_params - attributes["path"] = path - attributes["job_id"] = job_id - attributes["location"] = location - with span_creator.create('BigQuery.CancelJob', attributes): + + attributes = {"path": path, + "job_id": job_id, + "location": location + } + + with span_creator.create(name='BigQuery.CancelJob', attributes=attributes, client=self): resource = self._call_api( retry, method="POST", path=path, query_params=extra_params, timeout=timeout ) @@ -1780,11 +1773,12 @@ def list_jobs( path = "/projects/%s/jobs" % (project,) span_creator = SpanCreator() - attributes = extra_params - attributes["path"] = path - attributes["all_users"] = all_users - attributes["state_filter"] = state_filter - with span_creator.create('BigQuery.listJobs', attributes): + attributes = {"path": path, + "all_users": all_users, + "state_filter": state_filter, + } + + with span_creator.create(name='BigQuery.listJobs', attributes=attributes, client=self): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, @@ -2874,9 +2868,8 @@ def insert_rows_json( path = "%s/insertAll" % table.path # We can always retry, because every row has an insert ID. span_creator = SpanCreator() - attributes = data - attributes["path"] = path - with span_creator.create('BigQuery.insertRowsJson', attributes): + attributes = {"path": path} + with span_creator.create(name='BigQuery.insertRowsJson', attributes=attributes, client=self): response = self._call_api( retry, method="POST", @@ -3023,12 +3016,13 @@ def list_rows( params["startIndex"] = start_index span_creator = SpanCreator() path = "%s/data" % (table.path,) - attributes = params - attributes["path"] = path - attributes["page_size"] = page_size - attributes["page_token"] = page_token - attributes["start_index"] = start_index - with span_creator.create('BigQuery.listRows', attributes): + attributes = {"path": path, + "page_size": page_size, + "page_token": page_token, + "start_index": start_index, + } + + with span_creator.create(name='BigQuery.listRows', attributes=attributes, client=self): api_request = functools.partial(self._call_api, retry, timeout=timeout) row_iterator = RowIterator( client=self, diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index 0eec50506..c6927a566 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -638,7 +638,7 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None): # job has an ID. attributes = {'path': path} span_creator = SpanCreator() - with span_creator.create('BigQuery.job.begin', attributes): + with span_creator.create(name='BigQuery.job.begin', attributes=attributes, job_ref=self): api_response = client._call_api( retry, method="POST", path=path, data=self.to_api_repr(), timeout=timeout ) @@ -672,8 +672,8 @@ def exists(self, client=None, retry=DEFAULT_RETRY, timeout=None): try: span_creator = SpanCreator() attributes = extra_params - attributes['path'] = self.path - with span_creator.create('BigQuery.job.exists', attributes): + #attributes['path'] = self.path + with span_creator.create(name='BigQuery.job.exists', attributes=attributes, job_ref=self): client._call_api( retry, method="GET", @@ -709,8 +709,8 @@ def reload(self, client=None, retry=DEFAULT_RETRY, timeout=None): extra_params["location"] = self.location span_creator = SpanCreator() attributes = extra_params - attributes['path'] = self.path - with span_creator.create('BigQuery.job.reload', attributes): + #attributes['path'] = self.path + with span_creator.create(name='BigQuery.job.reload', attributes=attributes, job_ref=self): api_response = client._call_api( retry, method="GET", @@ -747,8 +747,8 @@ def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None): path = "{}/cancel".format(self.path) span_creator = SpanCreator() attributes = extra_params - attributes['path'] = path - with span_creator.create('BigQuery.job.cancel', attributes): + #attributes['path'] = path + with span_creator.create(name='BigQuery.job.cancel', attributes=attributes, job_ref=self): api_response = client._call_api( retry, method="POST", diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 280e61d1f..8e4c058d9 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -44,28 +44,59 @@ class SpanCreator: def __init__(self): + ''' + Constructs a span creator with all default attributes + ''' self.opentelemetry_enbled = HAS_OPENTELEMETRY + self.attributes = { + 'db.system': 'bigquery', + } + @contextmanager - def create(self, name, attributes=None): + def create(self, name, attributes=None, client=None, job_ref=None): if not self.opentelemetry_enbled: yield None return tracer = trace.get_tracer(__name__) - default_attributes = { - 'db.type': 'BigQuery', - } + if client: + self.set_client_attributes(client) + + elif job_ref: + self.set_job_attributes(job_ref) if attributes: - default_attributes.update(attributes) + self.attributes.update(attributes) # yield new span value - with tracer.start_as_current_span(name=name, attributes=default_attributes) as span: + with tracer.start_as_current_span(name=name, attributes=self.attributes) as span: try: yield span except GoogleAPICallError as error: if error.code is not None: span.set_status(Status(http_status_to_canonical_code(error.code))) raise + + def set_client_attributes(self, client): + self.attributes['db.name'] = client.project + self.attributes['location'] = client.location + + def set_job_attributes(self, job_ref): + self.attributes['db.name'] = job_ref.project + self.attributes['location'] = job_ref.location + self.attributes['num_child_jobs'] = str(job_ref.num_child_jobs) + self.attributes['job_id'] = job_ref.job_id + self.attributes['parent_job_id'] = job_ref.parent_job_id + #self.attributes['job_type'] = job_ref.job_type + self.attributes['timeCreated'] = job_ref.created + self.attributes['timeStarted'] = job_ref.started + self.attributes['timeEnded'] = job_ref.ended + self.attributes['errors'] = job_ref.errors + self.attributes['errorResult'] = job_ref.error_result + self.attributes['state'] = job_ref.state + + + + From 456b44588e98baec7f3b2da6973d1191c9e871b6 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Sat, 8 Aug 2020 23:50:14 -0400 Subject: [PATCH 07/41] adding tests, nox and correct attribute additions in client & job --- tests/unit/test_opentelemetry_tracing.py | 124 +++++++++++++++++++++++ 1 file changed, 124 insertions(+) create mode 100644 tests/unit/test_opentelemetry_tracing.py diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py new file mode 100644 index 000000000..015f06f19 --- /dev/null +++ b/tests/unit/test_opentelemetry_tracing.py @@ -0,0 +1,124 @@ +import sys +import pytest + +from importlib import reload +import unittest +import mock + +from google.cloud.bigquery import opentelemetry_tracing, client + +from opentelemetry import trace +from opentelemetry.sdk.trace import export +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor, ConsoleSpanExporter +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +import pdb + +TEST_SPAN_NAME = "bar" +TEST_SPAN_ATTRIBUTES = {"foo": "bar"} + + +# class TestOpentelemetry(unittest.TestCase): +@pytest.fixture +def setup(): + tracer_provider = TracerProvider() + memory_exporter = InMemorySpanExporter() + span_processor = SimpleExportSpanProcessor(memory_exporter) + tracer_provider.add_span_processor(span_processor) + trace.set_tracer_provider(tracer_provider) + yield memory_exporter + + +def test_opentelemetry_not_installed(setup): + temp_module = sys.modules["opentelemetry"] + sys.modules["opentelemetry"] = None + reload(opentelemetry_tracing) + span_creator = opentelemetry_tracing.SpanCreator() + with span_creator.create('No-op for opentelemetry') as span: + assert span == None + sys.modules["opentelemetry"] = temp_module + reload(opentelemetry_tracing) + + +def test_opentelemetry_success(setup): + span_creator = opentelemetry_tracing.SpanCreator() + expected_attributes = {"foo": "bar", + 'db.system': 'bigquery' + } + with span_creator.create( + TEST_SPAN_NAME, + attributes=TEST_SPAN_ATTRIBUTES, + ) as span: + if span is None: + span_list = setup.get_finished_spans() + print(span_list) + assert len(span_list) == 1 + assert span.name == TEST_SPAN_NAME + assert span.attributes == expected_attributes + + +def test_default_client_attributes(setup): + import google.auth.credentials + + mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) + test_client = client.Client(project='test_project', credentials=mock_credentials, + location='test_location') + + span_creator = opentelemetry_tracing.SpanCreator() + expected_attributes = {"foo": "bar", + 'db.system': 'bigquery', + 'db.name': 'test_project', + 'location': 'test_location' + } + with span_creator.create( + TEST_SPAN_NAME, + attributes=TEST_SPAN_ATTRIBUTES, + client=test_client + ) as span: + if span is None: + span_list = setup.get_finished_spans() + print(span_list) + assert len(span_list) == 1 + assert span.name == TEST_SPAN_NAME + assert span.attributes == expected_attributes + +def test_default_job_attributes(setup): + from google.cloud.bigquery import job + import google.auth.credentials + + mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) + + test_job_reference = job._JobReference(job_id='test_job_id', + project='test_project_id', + location='test_location') + test_client = client.Client(project='test_project', credentials=mock_credentials, + location='test_location') + test_job = job._AsyncJob(job_id=test_job_reference, client=test_client) + + span_creator = opentelemetry_tracing.SpanCreator() + # TODO add assertions for attributes being equal + expected_attributes = {'db.system': 'bigquery', + 'db.name': 'test_project_id', + 'location': 'test_location', + 'num_child_jobs': '0', + 'job_id': 'test_job_id', + 'parent_job_id': None, + 'timeCreated': None, + 'timeStarted': None, + 'timeEnded': None, + 'errors': None, + 'errorResult': None, + 'state': None, + 'foo': 'bar'} + + with span_creator.create( + TEST_SPAN_NAME, + attributes=TEST_SPAN_ATTRIBUTES, + job_ref=test_job + ) as span: + if span is None: + span_list = setup.get_finished_spans() + print(span_list[0]) + assert len(span_list) == 1 + assert span.name == TEST_SPAN_NAME + From f4aba6dac584f84a026c9e95c9c98a197c056a93 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Sat, 8 Aug 2020 23:51:35 -0400 Subject: [PATCH 08/41] adding tests, nox and correct attribute additions in client & job (left out of last commit) --- google/cloud/bigquery/client.py | 8 ++++---- google/cloud/bigquery/job.py | 9 +++------ .../cloud/bigquery/opentelemetry_tracing.py | 20 +++---------------- noxfile.py | 9 ++++++++- 4 files changed, 18 insertions(+), 28 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 34ee0ad7e..fd2f62233 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -1675,9 +1675,9 @@ def cancel_job( span_creator = SpanCreator() attributes = {"path": path, - "job_id": job_id, - "location": location - } + "job_id": job_id, + "location": location + } with span_creator.create(name='BigQuery.CancelJob', attributes=attributes, client=self): resource = self._call_api( @@ -2870,7 +2870,7 @@ def insert_rows_json( span_creator = SpanCreator() attributes = {"path": path} with span_creator.create(name='BigQuery.insertRowsJson', attributes=attributes, client=self): - response = self._call_api( + response = self._call_api( retry, method="POST", path=path, diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index c6927a566..c2a95362f 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -671,8 +671,7 @@ def exists(self, client=None, retry=DEFAULT_RETRY, timeout=None): try: span_creator = SpanCreator() - attributes = extra_params - #attributes['path'] = self.path + attributes = {'path': self.path} with span_creator.create(name='BigQuery.job.exists', attributes=attributes, job_ref=self): client._call_api( retry, @@ -708,8 +707,7 @@ def reload(self, client=None, retry=DEFAULT_RETRY, timeout=None): if self.location: extra_params["location"] = self.location span_creator = SpanCreator() - attributes = extra_params - #attributes['path'] = self.path + attributes = {'path': self.path} with span_creator.create(name='BigQuery.job.reload', attributes=attributes, job_ref=self): api_response = client._call_api( retry, @@ -746,8 +744,7 @@ def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None): path = "{}/cancel".format(self.path) span_creator = SpanCreator() - attributes = extra_params - #attributes['path'] = path + attributes = {'path': path} with span_creator.create(name='BigQuery.job.cancel', attributes=attributes, job_ref=self): api_response = client._call_api( retry, diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 8e4c058d9..9ba832998 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -14,18 +14,12 @@ import logging from contextlib import contextmanager - -from google.api_core.exceptions import GoogleAPICallError import pdb +from google.api_core.exceptions import GoogleAPICallError Logger = logging.getLogger(__name__) - try: from opentelemetry import trace - from opentelemetry import propagators - from opentelemetry.trace import SpanContext - from opentelemetry.trace import get_current_span - from opentelemetry.trace import set_span_in_context from opentelemetry.trace.status import Status from opentelemetry.instrumentation.utils import http_status_to_canonical_code @@ -44,15 +38,12 @@ class SpanCreator: def __init__(self): - ''' - Constructs a span creator with all default attributes - ''' + # Constructs a span creator with all default attributes self.opentelemetry_enbled = HAS_OPENTELEMETRY self.attributes = { 'db.system': 'bigquery', } - @contextmanager def create(self, name, attributes=None, client=None, job_ref=None): if not self.opentelemetry_enbled: @@ -69,7 +60,6 @@ def create(self, name, attributes=None, client=None, job_ref=None): if attributes: self.attributes.update(attributes) - # yield new span value with tracer.start_as_current_span(name=name, attributes=self.attributes) as span: try: @@ -89,14 +79,10 @@ def set_job_attributes(self, job_ref): self.attributes['num_child_jobs'] = str(job_ref.num_child_jobs) self.attributes['job_id'] = job_ref.job_id self.attributes['parent_job_id'] = job_ref.parent_job_id - #self.attributes['job_type'] = job_ref.job_type + # self.attributes['job_type'] = job_ref.job_type self.attributes['timeCreated'] = job_ref.created self.attributes['timeStarted'] = job_ref.started self.attributes['timeEnded'] = job_ref.ended self.attributes['errors'] = job_ref.errors self.attributes['errorResult'] = job_ref.error_result self.attributes['state'] = job_ref.state - - - - diff --git a/noxfile.py b/noxfile.py index bb6a10e1e..418038d2e 100644 --- a/noxfile.py +++ b/noxfile.py @@ -33,7 +33,14 @@ def default(session): """ # Install all test dependencies, then install local packages in-place. session.install( - "mock", "pytest", "google-cloud-testutils", "pytest-cov", "freezegun" + "mock", + "pytest", + "google-cloud-testutils", + "pytest-cov", + "freezegun", + "opentelemetry-api", + "opentelemetry-sdk", + "opentelemetry-instrumentation" ) session.install("grpcio") From e223657c7a23381ab51446f1c47000e452c0c66f Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Sun, 9 Aug 2020 13:03:44 -0400 Subject: [PATCH 09/41] linting --- google/cloud/bigquery/client.py | 568 ++++++++++-------- google/cloud/bigquery/job.py | 30 +- .../cloud/bigquery/opentelemetry_tracing.py | 42 +- tests/unit/test_opentelemetry_tracing.py | 78 +-- 4 files changed, 387 insertions(+), 331 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index fd2f62233..39b628aa1 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -167,14 +167,14 @@ class Client(ClientWithProject): """The scopes required for authenticating as a BigQuery consumer.""" def __init__( - self, - project=None, - credentials=None, - _http=None, - location=None, - default_query_job_config=None, - client_info=None, - client_options=None, + self, + project=None, + credentials=None, + _http=None, + location=None, + default_query_job_config=None, + client_info=None, + client_options=None, ): super(Client, self).__init__( project=project, credentials=credentials, _http=_http @@ -212,7 +212,7 @@ def close(self): self._http.close() def get_service_account_email( - self, project=None, retry=DEFAULT_RETRY, timeout=None + self, project=None, retry=DEFAULT_RETRY, timeout=None ): """Get the email address of the project's BigQuery service account @@ -244,13 +244,17 @@ def get_service_account_email( project = self.project path = "/projects/%s/serviceAccount" % (project,) span_creator = SpanCreator() - attributes = {'path': path} - with span_creator.create(name='BigQuery.getServiceAccountEmail', attributes=attributes, client=self): - api_response = self._call_api(retry, method="GET", path=path, timeout=timeout) + attributes = {"path": path} + with span_creator.create( + name="BigQuery.getServiceAccountEmail", attributes=attributes, client=self + ): + api_response = self._call_api( + retry, method="GET", path=path, timeout=timeout + ) return api_response["email"] def list_projects( - self, max_results=None, page_token=None, retry=DEFAULT_RETRY, timeout=None + self, max_results=None, page_token=None, retry=DEFAULT_RETRY, timeout=None ): """List projects for the project associated with this client. @@ -280,10 +284,12 @@ def list_projects( Iterator of :class:`~google.cloud.bigquery.client.Project` accessible to the current client. """ - path = '/projects' + path = "/projects" span_creator = SpanCreator() - attributes = {'path': path} - with span_creator.create(name='BigQuery.list_projects', attributes=attributes, client=self): + attributes = {"path": path} + with span_creator.create( + name="BigQuery.list_projects", attributes=attributes, client=self + ): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, @@ -296,14 +302,14 @@ def list_projects( ) def list_datasets( - self, - project=None, - include_all=False, - filter=None, - max_results=None, - page_token=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + project=None, + include_all=False, + filter=None, + max_results=None, + page_token=None, + retry=DEFAULT_RETRY, + timeout=None, ): """List datasets for the project associated with this client. @@ -350,11 +356,11 @@ def list_datasets( extra_params["filter"] = filter path = "/projects/%s/datasets" % (project,) span_creator = SpanCreator() - attributes = {'path': path, - 'page_token': page_token - } + attributes = {"path": path, "page_token": page_token} - with span_creator.create(name='BigQuery.list_datasets', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.list_datasets", attributes=attributes, client=self + ): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, @@ -427,7 +433,7 @@ def _create_bqstorage_client(self): return bigquery_storage_v1.BigQueryReadClient(credentials=self._credentials) def create_dataset( - self, dataset, exists_ok=False, retry=DEFAULT_RETRY, timeout=None + self, dataset, exists_ok=False, retry=DEFAULT_RETRY, timeout=None ): """API call: create the dataset via a POST request. @@ -483,8 +489,10 @@ def create_dataset( try: span_creator = SpanCreator() - attributes = {'path': path} - with span_creator.create(name='BigQuery.createDataset', attributes=attributes, client=self): + attributes = {"path": path} + with span_creator.create( + name="BigQuery.createDataset", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="POST", path=path, data=data, timeout=timeout ) @@ -495,7 +503,7 @@ def create_dataset( return self.get_dataset(dataset.reference, retry=retry) def create_routine( - self, routine, exists_ok=False, retry=DEFAULT_RETRY, timeout=None + self, routine, exists_ok=False, retry=DEFAULT_RETRY, timeout=None ): """[Beta] Create a routine via a POST request. @@ -530,9 +538,10 @@ def create_routine( resource = routine.to_api_repr() try: span_creator = SpanCreator() - attributes = {'path': path, - 'exists_ok': exists_ok} - with span_creator.create(name='BigQuery.createRoutine', attributes=attributes, client=self): + attributes = {"path": path, "exists_ok": exists_ok} + with span_creator.create( + name="BigQuery.createRoutine", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="POST", path=path, data=resource, timeout=timeout ) @@ -581,9 +590,10 @@ def create_table(self, table, exists_ok=False, retry=DEFAULT_RETRY, timeout=None data = table.to_api_repr() try: span_creator = SpanCreator() - attributes = {'path': path, - 'dataset_id': dataset_id} - with span_creator.create(name='BigQuery.createTable', attributes=attributes, client=self): + attributes = {"path": path, "dataset_id": dataset_id} + with span_creator.create( + name="BigQuery.createTable", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="POST", path=path, data=data, timeout=timeout ) @@ -628,15 +638,17 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None): ) span_creator = SpanCreator() path = dataset_ref.path - attributes = {'path': path} - with span_creator.create(name='BigQuery.createTable', attributes=attributes, client=self): + attributes = {"path": path} + with span_creator.create( + name="BigQuery.createTable", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout ) return Dataset.from_api_repr(api_response) def get_iam_policy( - self, table, requested_policy_version=1, retry=DEFAULT_RETRY, timeout=None, + self, table, requested_policy_version=1, retry=DEFAULT_RETRY, timeout=None, ): if not isinstance(table, (Table, TableReference)): raise TypeError("table must be a Table or TableReference") @@ -648,9 +660,10 @@ def get_iam_policy( path = "{}:getIamPolicy".format(table.path) span_creator = SpanCreator() - attributes = {'path': path, - 'body': body} - with span_creator.create(name='BigQuery.getIAMPolicy', attributes=attributes, client=self): + attributes = {"path": path, "body": body} + with span_creator.create( + name="BigQuery.getIAMPolicy", attributes=attributes, client=self + ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, ) @@ -658,7 +671,7 @@ def get_iam_policy( return Policy.from_api_repr(response) def set_iam_policy( - self, table, policy, updateMask=None, retry=DEFAULT_RETRY, timeout=None, + self, table, policy, updateMask=None, retry=DEFAULT_RETRY, timeout=None, ): if not isinstance(table, (Table, TableReference)): raise TypeError("table must be a Table or TableReference") @@ -673,9 +686,10 @@ def set_iam_policy( path = "{}:setIamPolicy".format(table.path) span_creator = SpanCreator() - attributes = {'path': path, - 'body': body} - with span_creator.create(name='BigQuery.setIAMPolicy', attributes=attributes, client=self): + attributes = {"path": path, "body": body} + with span_creator.create( + name="BigQuery.setIAMPolicy", attributes=attributes, client=self + ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, ) @@ -683,7 +697,7 @@ def set_iam_policy( return Policy.from_api_repr(response) def test_iam_permissions( - self, table, permissions, retry=DEFAULT_RETRY, timeout=None, + self, table, permissions, retry=DEFAULT_RETRY, timeout=None, ): if not isinstance(table, (Table, TableReference)): raise TypeError("table must be a Table or TableReference") @@ -692,9 +706,10 @@ def test_iam_permissions( path = "{}:testIamPermissions".format(table.path) span_creator = SpanCreator() - attributes = {'path': path, - 'body': body} - with span_creator.create(name='BigQuery.testIAMPolicy', attributes=attributes, client=self): + attributes = {"path": path, "body": body} + with span_creator.create( + name="BigQuery.testIAMPolicy", attributes=attributes, client=self + ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, ) @@ -728,8 +743,10 @@ def get_model(self, model_ref, retry=DEFAULT_RETRY, timeout=None): ) span_creator = SpanCreator() path = model_ref.path - attributes = {'path': path} - with span_creator.create(name='BigQuery.getModel', attributes=attributes, client=self): + attributes = {"path": path} + with span_creator.create( + name="BigQuery.getModel", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout ) @@ -765,8 +782,10 @@ def get_routine(self, routine_ref, retry=DEFAULT_RETRY, timeout=None): ) path = routine_ref.path span_creator = SpanCreator() - attributes = {'path': path} - with span_creator.create(name='BigQuery.getRoutine', attributes=attributes, client=self): + attributes = {"path": path} + with span_creator.create( + name="BigQuery.getRoutine", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout ) @@ -798,8 +817,10 @@ 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_creator = SpanCreator() - attributes = {'path': path} - with span_creator.create(name='BigQuery.getTable', attributes=attributes, client=self): + attributes = {"path": path} + with span_creator.create( + name="BigQuery.getTable", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout ) @@ -841,10 +862,11 @@ def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): headers = None path = dataset.path span_creator = SpanCreator() - attributes = {'path': path, - 'fields': fields} + attributes = {"path": path, "fields": fields} - with span_creator.create(name='BigQuery.updateDataset', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.updateDataset", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="PATCH", @@ -890,10 +912,11 @@ def update_model(self, model, fields, retry=DEFAULT_RETRY, timeout=None): headers = None path = model.path span_creator = SpanCreator() - attributes = {'path': path, - 'fields': fields} + attributes = {"path": path, "fields": fields} - with span_creator.create(name='BigQuery.updateModel', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.updateModel", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="PATCH", @@ -950,10 +973,11 @@ def update_routine(self, routine, fields, retry=DEFAULT_RETRY, timeout=None): path = routine.path span_creator = SpanCreator() - attributes = {'path': path, - 'fields': fields} + attributes = {"path": path, "fields": fields} - with span_creator.create(name='BigQuery.updateRountine', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.updateRountine", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="PUT", @@ -1000,10 +1024,11 @@ def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): path = table.path span_creator = SpanCreator() - attributes = {'path': path, - 'fields': fields} + attributes = {"path": path, "fields": fields} - with span_creator.create(name='BigQuery.updateTable', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.updateTable", attributes=attributes, client=self + ): api_response = self._call_api( retry, method="PATCH", @@ -1015,12 +1040,12 @@ def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): return Table.from_api_repr(api_response) def list_models( - self, - dataset, - max_results=None, - page_token=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + dataset, + max_results=None, + page_token=None, + retry=DEFAULT_RETRY, + timeout=None, ): """[Beta] List models in the dataset. @@ -1068,10 +1093,14 @@ def list_models( path = "%s/models" % dataset.path span_creator = SpanCreator() - attributes = {'path': path, - 'page_token': page_token, - 'max_results': max_results} - with span_creator.create(name='BigQuery.listModels', attributes=attributes, client=self): + attributes = { + "path": path, + "page_token": page_token, + "max_results": max_results, + } + with span_creator.create( + name="BigQuery.listModels", attributes=attributes, client=self + ): api_request = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( client=self, @@ -1086,12 +1115,12 @@ def list_models( return result def list_routines( - self, - dataset, - max_results=None, - page_token=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + dataset, + max_results=None, + page_token=None, + retry=DEFAULT_RETRY, + timeout=None, ): """[Beta] List routines in the dataset. @@ -1139,10 +1168,14 @@ def list_routines( path = "{}/routines".format(dataset.path) span_creator = SpanCreator() - attributes = {'path': path, - 'page_token': page_token, - 'max_results': max_results} - with span_creator.create(name='BigQuery.listRoutines', attributes=attributes, client=self): + attributes = { + "path": path, + "page_token": page_token, + "max_results": max_results, + } + with span_creator.create( + name="BigQuery.listRoutines", attributes=attributes, client=self + ): api_requests = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( client=self, @@ -1157,12 +1190,12 @@ def list_routines( return result def list_tables( - self, - dataset, - max_results=None, - page_token=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + dataset, + max_results=None, + page_token=None, + retry=DEFAULT_RETRY, + timeout=None, ): """List tables in the dataset. @@ -1210,10 +1243,14 @@ def list_tables( path = "%s/tables" % dataset.path span_creator = SpanCreator() - attributes = {'path': path, - 'page_token': page_token, - 'max_results': max_results} - with span_creator.create(name='BigQuery.listTables', attributes=attributes, client=self): + attributes = { + "path": path, + "page_token": page_token, + "max_results": max_results, + } + with span_creator.create( + name="BigQuery.listTables", attributes=attributes, client=self + ): api_requests = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( @@ -1229,12 +1266,12 @@ def list_tables( return result def delete_dataset( - self, - dataset, - delete_contents=False, - retry=DEFAULT_RETRY, - timeout=None, - not_found_ok=False, + self, + dataset, + delete_contents=False, + retry=DEFAULT_RETRY, + timeout=None, + not_found_ok=False, ): """Delete a dataset. @@ -1280,7 +1317,9 @@ def delete_dataset( try: span_creator = SpanCreator() attributes = {"path": path} - with span_creator.create(name='BigQuery.deleteDataset', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.deleteDataset", attributes=attributes, client=self + ): self._call_api( retry, method="DELETE", @@ -1293,7 +1332,7 @@ def delete_dataset( raise def delete_model( - self, model, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False + self, model, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False ): """[Beta] Delete a model @@ -1328,16 +1367,17 @@ def delete_model( path = model.path try: span_creator = SpanCreator() - attributes = {"not_found_okay": not_found_ok, - "path": path} - with span_creator.create(name='BigQuery.deleteModel', attributes=attributes, client=self): + attributes = {"not_found_okay": not_found_ok, "path": path} + with span_creator.create( + name="BigQuery.deleteModel", attributes=attributes, client=self + ): self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: raise def delete_routine( - self, routine, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False + self, routine, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False ): """[Beta] Delete a routine. @@ -1375,14 +1415,16 @@ def delete_routine( try: span_creator = SpanCreator() attributes = {"path": path} - with span_creator.create(name='BigQuery.deleteRoutine', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.deleteRoutine", attributes=attributes, client=self + ): self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: raise def delete_table( - self, table, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False + self, table, retry=DEFAULT_RETRY, timeout=None, not_found_ok=False ): """Delete a table @@ -1416,14 +1458,16 @@ def delete_table( path = table.path span_creator = SpanCreator() attributes = {"path": path} - with span_creator.create(name='BigQuery.deleteTable', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.deleteTable", attributes=attributes, client=self + ): self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: if not not_found_ok: raise def _get_query_results( - self, job_id, retry, project=None, timeout_ms=None, location=None, timeout=None + self, job_id, retry, project=None, timeout_ms=None, location=None, timeout=None ): """Get the query results object for a query job. @@ -1466,13 +1510,17 @@ def _get_query_results( # job is complete (from QueryJob.done(), called ultimately from # QueryJob.result()). So we don't need to poll here. span_creator = SpanCreator() - attributes = {"path": path, - "job_id": job_id, - "location": location} + attributes = {"path": path, "job_id": job_id, "location": location} - with span_creator.create(name='BigQuery.getQueryResults', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.getQueryResults", attributes=attributes, client=self + ): resource = self._call_api( - retry, method="GET", path=path, query_params=extra_params, timeout=timeout + retry, + method="GET", + path=path, + query_params=extra_params, + timeout=timeout, ) return _QueryResults.from_api_repr(resource) @@ -1575,7 +1623,7 @@ def create_job(self, job_config, retry=DEFAULT_RETRY): raise TypeError("Invalid job configuration received.") def get_job( - self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None + self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None ): """Fetch a job for the project associated with this client. @@ -1618,20 +1666,23 @@ def get_job( path = "/projects/{}/jobs/{}".format(project, job_id) span_creator = SpanCreator() - attributes = {"path": path, - "job_id": job_id, - "location": location - } + attributes = {"path": path, "job_id": job_id, "location": location} - with span_creator.create(name='BigQuery.getJob', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.getJob", attributes=attributes, client=self + ): resource = self._call_api( - retry, method="GET", path=path, query_params=extra_params, timeout=timeout + retry, + method="GET", + path=path, + query_params=extra_params, + timeout=timeout, ) return self.job_from_resource(resource) def cancel_job( - self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None + self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None ): """Attempt to cancel a job from a job ID. @@ -1674,30 +1725,33 @@ def cancel_job( path = "/projects/{}/jobs/{}/cancel".format(project, job_id) span_creator = SpanCreator() - attributes = {"path": path, - "job_id": job_id, - "location": location - } + attributes = {"path": path, "job_id": job_id, "location": location} - with span_creator.create(name='BigQuery.CancelJob', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.CancelJob", attributes=attributes, client=self + ): resource = self._call_api( - retry, method="POST", path=path, query_params=extra_params, timeout=timeout + retry, + method="POST", + path=path, + query_params=extra_params, + timeout=timeout, ) return self.job_from_resource(resource["job"]) def list_jobs( - self, - project=None, - parent_job=None, - max_results=None, - page_token=None, - all_users=None, - state_filter=None, - retry=DEFAULT_RETRY, - timeout=None, - min_creation_time=None, - max_creation_time=None, + self, + project=None, + parent_job=None, + max_results=None, + page_token=None, + all_users=None, + state_filter=None, + retry=DEFAULT_RETRY, + timeout=None, + min_creation_time=None, + max_creation_time=None, ): """List jobs for the project associated with this client. @@ -1773,12 +1827,15 @@ def list_jobs( path = "/projects/%s/jobs" % (project,) span_creator = SpanCreator() - attributes = {"path": path, - "all_users": all_users, - "state_filter": state_filter, - } + attributes = { + "path": path, + "all_users": all_users, + "state_filter": state_filter, + } - with span_creator.create(name='BigQuery.listJobs', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.listJobs", attributes=attributes, client=self + ): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, @@ -1792,16 +1849,16 @@ def list_jobs( ) def load_table_from_uri( - self, - source_uris, - destination, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + source_uris, + destination, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, + retry=DEFAULT_RETRY, + timeout=None, ): """Starts a job for loading data into a table from CloudStorage. @@ -1874,17 +1931,17 @@ def load_table_from_uri( return load_job def load_table_from_file( - self, - file_obj, - destination, - rewind=False, - size=None, - num_retries=_DEFAULT_NUM_RETRIES, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, + self, + file_obj, + destination, + rewind=False, + size=None, + num_retries=_DEFAULT_NUM_RETRIES, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, ): """Upload the contents of this table from a file-like object. @@ -1974,16 +2031,16 @@ def load_table_from_file( return self.job_from_resource(response.json()) def load_table_from_dataframe( - self, - dataframe, - destination, - num_retries=_DEFAULT_NUM_RETRIES, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, - parquet_compression="snappy", + self, + dataframe, + destination, + num_retries=_DEFAULT_NUM_RETRIES, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, + parquet_compression="snappy", ): """Upload the contents of a table from a pandas DataFrame. @@ -2082,8 +2139,8 @@ def load_table_from_dataframe( # schema, and check if dataframe schema is compatible with it - except # for WRITE_TRUNCATE jobs, the existing schema does not matter then. if ( - not job_config.schema - and job_config.write_disposition != job.WriteDisposition.WRITE_TRUNCATE + not job_config.schema + and job_config.write_disposition != job.WriteDisposition.WRITE_TRUNCATE ): try: table = self.get_table(destination) @@ -2165,15 +2222,15 @@ def load_table_from_dataframe( os.remove(tmppath) def load_table_from_json( - self, - json_rows, - destination, - num_retries=_DEFAULT_NUM_RETRIES, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, + self, + json_rows, + destination, + num_retries=_DEFAULT_NUM_RETRIES, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, ): """Upload the contents of a table from a JSON string or dict. @@ -2380,16 +2437,16 @@ def _do_multipart_upload(self, stream, metadata, size, num_retries): return response def copy_table( - self, - sources, - destination, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + sources, + destination, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, + retry=DEFAULT_RETRY, + timeout=None, ): """Copy one or more tables to another table. @@ -2483,17 +2540,17 @@ def copy_table( return copy_job def extract_table( - self, - source, - destination_uris, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - job_config=None, - retry=DEFAULT_RETRY, - timeout=None, - source_type="Table", + self, + source, + destination_uris, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + job_config=None, + retry=DEFAULT_RETRY, + timeout=None, + source_type="Table", ): """Start a job to extract a table into Cloud Storage files. @@ -2582,15 +2639,15 @@ def extract_table( return extract_job def query( - self, - query, - job_config=None, - job_id=None, - job_id_prefix=None, - location=None, - project=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + query, + job_config=None, + job_id=None, + job_id_prefix=None, + location=None, + project=None, + retry=DEFAULT_RETRY, + timeout=None, ): """Run a SQL query. @@ -2731,7 +2788,7 @@ def insert_rows(self, table, rows, selected_fields=None, **kwargs): return self.insert_rows_json(table, json_rows, **kwargs) def insert_rows_from_dataframe( - self, table, dataframe, selected_fields=None, chunk_size=500, **kwargs + self, table, dataframe, selected_fields=None, chunk_size=500, **kwargs ): """Insert rows into a table from a dataframe via the streaming API. @@ -2779,15 +2836,15 @@ def insert_rows_from_dataframe( return insert_results def insert_rows_json( - self, - table, - json_rows, - row_ids=None, - skip_invalid_rows=None, - ignore_unknown_values=None, - template_suffix=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + table, + json_rows, + row_ids=None, + skip_invalid_rows=None, + ignore_unknown_values=None, + template_suffix=None, + retry=DEFAULT_RETRY, + timeout=None, ): """Insert rows into a table without applying local type conversions. @@ -2838,7 +2895,7 @@ def insert_rows_json( TypeError: if `json_rows` is not a `Sequence`. """ if not isinstance( - json_rows, (collections_abc.Sequence, collections_abc.Iterator) + json_rows, (collections_abc.Sequence, collections_abc.Iterator) ): raise TypeError("json_rows argument should be a sequence of dicts") # Convert table to just a reference because unlike insert_rows, @@ -2869,13 +2926,11 @@ def insert_rows_json( # We can always retry, because every row has an insert ID. span_creator = SpanCreator() attributes = {"path": path} - with span_creator.create(name='BigQuery.insertRowsJson', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.insertRowsJson", attributes=attributes, client=self + ): response = self._call_api( - retry, - method="POST", - path=path, - data=data, - timeout=timeout, + retry, method="POST", path=path, data=data, timeout=timeout, ) errors = [] @@ -2925,15 +2980,15 @@ def list_partitions(self, table, retry=DEFAULT_RETRY, timeout=None): ] def list_rows( - self, - table, - selected_fields=None, - max_results=None, - page_token=None, - start_index=None, - page_size=None, - retry=DEFAULT_RETRY, - timeout=None, + self, + table, + selected_fields=None, + max_results=None, + page_token=None, + start_index=None, + page_size=None, + retry=DEFAULT_RETRY, + timeout=None, ): """List the rows of the table. @@ -3016,13 +3071,16 @@ def list_rows( params["startIndex"] = start_index span_creator = SpanCreator() path = "%s/data" % (table.path,) - attributes = {"path": path, - "page_size": page_size, - "page_token": page_token, - "start_index": start_index, - } + attributes = { + "path": path, + "page_size": page_size, + "page_token": page_token, + "start_index": start_index, + } - with span_creator.create(name='BigQuery.listRows', attributes=attributes, client=self): + with span_creator.create( + name="BigQuery.listRows", attributes=attributes, client=self + ): api_request = functools.partial(self._call_api, retry, timeout=timeout) row_iterator = RowIterator( client=self, diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index c2a95362f..5ee86cb39 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -636,11 +636,17 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None): # jobs.insert is idempotent because we ensure that every new # job has an ID. - attributes = {'path': path} + attributes = {"path": path} span_creator = SpanCreator() - with span_creator.create(name='BigQuery.job.begin', attributes=attributes, job_ref=self): + with span_creator.create( + name="BigQuery.job.begin", attributes=attributes, job_ref=self + ): api_response = client._call_api( - retry, method="POST", path=path, data=self.to_api_repr(), timeout=timeout + retry, + method="POST", + path=path, + data=self.to_api_repr(), + timeout=timeout, ) self._set_properties(api_response) @@ -671,8 +677,10 @@ def exists(self, client=None, retry=DEFAULT_RETRY, timeout=None): try: span_creator = SpanCreator() - attributes = {'path': self.path} - with span_creator.create(name='BigQuery.job.exists', attributes=attributes, job_ref=self): + attributes = {"path": self.path} + with span_creator.create( + name="BigQuery.job.exists", attributes=attributes, job_ref=self + ): client._call_api( retry, method="GET", @@ -707,8 +715,10 @@ def reload(self, client=None, retry=DEFAULT_RETRY, timeout=None): if self.location: extra_params["location"] = self.location span_creator = SpanCreator() - attributes = {'path': self.path} - with span_creator.create(name='BigQuery.job.reload', attributes=attributes, job_ref=self): + attributes = {"path": self.path} + with span_creator.create( + name="BigQuery.job.reload", attributes=attributes, job_ref=self + ): api_response = client._call_api( retry, method="GET", @@ -744,8 +754,10 @@ def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None): path = "{}/cancel".format(self.path) span_creator = SpanCreator() - attributes = {'path': path} - with span_creator.create(name='BigQuery.job.cancel', attributes=attributes, job_ref=self): + attributes = {"path": path} + with span_creator.create( + name="BigQuery.job.cancel", attributes=attributes, job_ref=self + ): api_response = client._call_api( retry, method="POST", diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 9ba832998..e2d8f483f 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -14,7 +14,6 @@ import logging from contextlib import contextmanager -import pdb from google.api_core.exceptions import GoogleAPICallError Logger = logging.getLogger(__name__) @@ -27,10 +26,10 @@ except ImportError: Logger.info( - 'This service instrumented using opentelemetry.' - 'Opentelemetry could not be imported please' - 'add opentelemetry-api and opentelemetry-instrumentation' - 'packages in order to get Big Query Tracing data' + "This service instrumented using opentelemetry." + "Opentelemetry could not be imported please" + "add opentelemetry-api and opentelemetry-instrumentation" + "packages in order to get Big Query Tracing data" ) HAS_OPENTELEMETRY = False @@ -41,7 +40,7 @@ def __init__(self): # Constructs a span creator with all default attributes self.opentelemetry_enbled = HAS_OPENTELEMETRY self.attributes = { - 'db.system': 'bigquery', + "db.system": "bigquery", } @contextmanager @@ -61,7 +60,9 @@ def create(self, name, attributes=None, client=None, job_ref=None): if attributes: self.attributes.update(attributes) # yield new span value - with tracer.start_as_current_span(name=name, attributes=self.attributes) as span: + with tracer.start_as_current_span( + name=name, attributes=self.attributes + ) as span: try: yield span except GoogleAPICallError as error: @@ -70,19 +71,18 @@ def create(self, name, attributes=None, client=None, job_ref=None): raise def set_client_attributes(self, client): - self.attributes['db.name'] = client.project - self.attributes['location'] = client.location + self.attributes["db.name"] = client.project + self.attributes["location"] = client.location def set_job_attributes(self, job_ref): - self.attributes['db.name'] = job_ref.project - self.attributes['location'] = job_ref.location - self.attributes['num_child_jobs'] = str(job_ref.num_child_jobs) - self.attributes['job_id'] = job_ref.job_id - self.attributes['parent_job_id'] = job_ref.parent_job_id - # self.attributes['job_type'] = job_ref.job_type - self.attributes['timeCreated'] = job_ref.created - self.attributes['timeStarted'] = job_ref.started - self.attributes['timeEnded'] = job_ref.ended - self.attributes['errors'] = job_ref.errors - self.attributes['errorResult'] = job_ref.error_result - self.attributes['state'] = job_ref.state + self.attributes["db.name"] = job_ref.project + self.attributes["location"] = job_ref.location + self.attributes["num_child_jobs"] = str(job_ref.num_child_jobs) + self.attributes["job_id"] = job_ref.job_id + self.attributes["parent_job_id"] = job_ref.parent_job_id + self.attributes["timeCreated"] = job_ref.created + self.attributes["timeStarted"] = job_ref.started + self.attributes["timeEnded"] = job_ref.ended + self.attributes["errors"] = job_ref.errors + self.attributes["errorResult"] = job_ref.error_result + self.attributes["state"] = job_ref.state diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 015f06f19..eb7630eb4 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -2,23 +2,19 @@ import pytest from importlib import reload -import unittest import mock from google.cloud.bigquery import opentelemetry_tracing, client from opentelemetry import trace -from opentelemetry.sdk.trace import export from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor, ConsoleSpanExporter +from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -import pdb TEST_SPAN_NAME = "bar" TEST_SPAN_ATTRIBUTES = {"foo": "bar"} -# class TestOpentelemetry(unittest.TestCase): @pytest.fixture def setup(): tracer_provider = TracerProvider() @@ -34,21 +30,16 @@ def test_opentelemetry_not_installed(setup): sys.modules["opentelemetry"] = None reload(opentelemetry_tracing) span_creator = opentelemetry_tracing.SpanCreator() - with span_creator.create('No-op for opentelemetry') as span: - assert span == None + with span_creator.create("No-op for opentelemetry") as span: + assert span is None sys.modules["opentelemetry"] = temp_module reload(opentelemetry_tracing) def test_opentelemetry_success(setup): span_creator = opentelemetry_tracing.SpanCreator() - expected_attributes = {"foo": "bar", - 'db.system': 'bigquery' - } - with span_creator.create( - TEST_SPAN_NAME, - attributes=TEST_SPAN_ATTRIBUTES, - ) as span: + expected_attributes = {"foo": "bar", "db.system": "bigquery"} + with span_creator.create(TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES,) as span: if span is None: span_list = setup.get_finished_spans() print(span_list) @@ -61,19 +52,19 @@ def test_default_client_attributes(setup): import google.auth.credentials mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) - test_client = client.Client(project='test_project', credentials=mock_credentials, - location='test_location') + test_client = client.Client( + project="test_project", credentials=mock_credentials, location="test_location" + ) span_creator = opentelemetry_tracing.SpanCreator() - expected_attributes = {"foo": "bar", - 'db.system': 'bigquery', - 'db.name': 'test_project', - 'location': 'test_location' - } + expected_attributes = { + "foo": "bar", + "db.system": "bigquery", + "db.name": "test_project", + "location": "test_location", + } with span_creator.create( - TEST_SPAN_NAME, - attributes=TEST_SPAN_ATTRIBUTES, - client=test_client + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client ) as span: if span is None: span_list = setup.get_finished_spans() @@ -82,43 +73,38 @@ def test_default_client_attributes(setup): assert span.name == TEST_SPAN_NAME assert span.attributes == expected_attributes + def test_default_job_attributes(setup): from google.cloud.bigquery import job import google.auth.credentials mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) - test_job_reference = job._JobReference(job_id='test_job_id', - project='test_project_id', - location='test_location') - test_client = client.Client(project='test_project', credentials=mock_credentials, - location='test_location') + test_job_reference = job._JobReference( + job_id="test_job_id", project="test_project_id", location="test_location" + ) + test_client = client.Client( + project="test_project", credentials=mock_credentials, location="test_location" + ) test_job = job._AsyncJob(job_id=test_job_reference, client=test_client) span_creator = opentelemetry_tracing.SpanCreator() # TODO add assertions for attributes being equal - expected_attributes = {'db.system': 'bigquery', - 'db.name': 'test_project_id', - 'location': 'test_location', - 'num_child_jobs': '0', - 'job_id': 'test_job_id', - 'parent_job_id': None, - 'timeCreated': None, - 'timeStarted': None, - 'timeEnded': None, - 'errors': None, - 'errorResult': None, - 'state': None, - 'foo': 'bar'} + expected_attributes = { + "db.system": "bigquery", + "db.name": "test_project_id", + "location": "test_location", + "num_child_jobs": "0", + "job_id": "test_job_id", + "foo": "bar", + } with span_creator.create( - TEST_SPAN_NAME, - attributes=TEST_SPAN_ATTRIBUTES, - job_ref=test_job + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job ) as span: if span is None: span_list = setup.get_finished_spans() print(span_list[0]) assert len(span_list) == 1 assert span.name == TEST_SPAN_NAME - + assert span.attributes == expected_attributes From 78b82980f5f9db6b60542e7620e4f1639ce4e61f Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Sun, 9 Aug 2020 13:22:17 -0400 Subject: [PATCH 10/41] reformatting noxfile.[y --- noxfile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index 418038d2e..8615f00db 100644 --- a/noxfile.py +++ b/noxfile.py @@ -40,7 +40,7 @@ def default(session): "freezegun", "opentelemetry-api", "opentelemetry-sdk", - "opentelemetry-instrumentation" + "opentelemetry-instrumentation", ) session.install("grpcio") From 95a35ad0eef64fc53d0f4d9464237a1504766ff0 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Tue, 11 Aug 2020 16:08:59 -0400 Subject: [PATCH 11/41] addressing suggested changes --- README.rst | 29 +-- google/cloud/bigquery/client.py | 214 ++++++++---------- google/cloud/bigquery/job.py | 30 ++- .../cloud/bigquery/opentelemetry_tracing.py | 107 +++++---- setup.py | 6 + tests/unit/test_opentelemetry_tracing.py | 21 +- 6 files changed, 184 insertions(+), 223 deletions(-) diff --git a/README.rst b/README.rst index 7141b8456..145b8f1c8 100644 --- a/README.rst +++ b/README.rst @@ -106,28 +106,19 @@ Perform a query Instrumenting With OpenTelemetry -------------------------------- -This application uses OpenTelmetry to output tracing data from -api calls to BigQuery. To enable OpenTelemetry tracing in -the BigQuery client the following PyPi packages need to be installed: +This application uses `OpenTelemetry`_ +to output tracing data from +API calls to BigQuery. To enable OpenTelemetry tracing in +the BigQuery client the following PyPI packages need to be installed: -.. code-block:: console - - pip install opentelemetry-api - -.. code-block:: console - - pip install opentelemetry-sdk - -.. code-block:: console - - pip install opentelemetry-instrumentaion +.. _OpenTelemetry: https://opentelemetry.io .. code-block:: console - pip install opentelemetry-exporter-cloud_trace + pip install google-cloud-bigquery[opentelemetry] -After installation, OpenTelemetry can now be used in the BigQuery -client and in BigQuery jobs. First however an exporter must be +After installation, OpenTelemetry can be used in the BigQuery +client and in BigQuery jobs. First, however, an exporter must be specified for where the trace data will be outputted to. An example of this can be found here: @@ -135,11 +126,11 @@ can be found here: from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor + from opentelemetry.sdk.trace.export import BatchExportSpanProcessor from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter trace.set_tracer_provider(TracerProvider()) trace.get_tracer_provider().add_span_processor( - SimpleExportSpanProcessor(CloudTraceSpanExporter()) + BatchExportSpanProcessor(CloudTraceSpanExporter()) ) In this example all tracing data will be published to the Google diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 39b628aa1..ad28cb87e 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -33,7 +33,7 @@ import tempfile import uuid import warnings -from google.cloud.bigquery.opentelemetry_tracing import SpanCreator +from google.cloud.bigquery.opentelemetry_tracing import create_span try: import pyarrow @@ -243,10 +243,11 @@ def get_service_account_email( if project is None: project = self.project path = "/projects/%s/serviceAccount" % (project,) - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.getServiceAccountEmail", attributes=attributes, client=self + 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 @@ -285,10 +286,9 @@ def list_projects( accessible to the current client. """ path = "/projects" - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.list_projects", attributes=attributes, client=self + span_attributes = {"path": path} + with create_span( + name="BigQuery.list_projects", attributes=span_attributes, client=self ): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( @@ -355,11 +355,10 @@ def list_datasets( # and converting it into a string here. extra_params["filter"] = filter path = "/projects/%s/datasets" % (project,) - span_creator = SpanCreator() - attributes = {"path": path, "page_token": page_token} + span_attributes = {"path": path, "page_token": page_token} - with span_creator.create( - name="BigQuery.list_datasets", attributes=attributes, client=self + with create_span( + name="BigQuery.list_datasets", attributes=span_attributes, client=self ): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( @@ -488,10 +487,9 @@ def create_dataset( data["location"] = self.location try: - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.createDataset", attributes=attributes, client=self + 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 @@ -537,10 +535,9 @@ def create_routine( ) resource = routine.to_api_repr() try: - span_creator = SpanCreator() - attributes = {"path": path, "exists_ok": exists_ok} - with span_creator.create( - name="BigQuery.createRoutine", attributes=attributes, client=self + span_attributes = {"path": path, "exists_ok": exists_ok} + 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 @@ -589,10 +586,9 @@ def create_table(self, table, exists_ok=False, retry=DEFAULT_RETRY, timeout=None path = "/projects/%s/datasets/%s/tables" % (table.project, dataset_id) data = table.to_api_repr() try: - span_creator = SpanCreator() - attributes = {"path": path, "dataset_id": dataset_id} - with span_creator.create( - name="BigQuery.createTable", attributes=attributes, client=self + 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 @@ -636,11 +632,10 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None): dataset_ref = DatasetReference.from_string( dataset_ref, default_project=self.project ) - span_creator = SpanCreator() path = dataset_ref.path - attributes = {"path": path} - with span_creator.create( - name="BigQuery.createTable", attributes=attributes, client=self + span_attributes = {"path": path} + with create_span( + name="BigQuery.createTable", attributes=span_attributes, client=self ): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout @@ -659,10 +654,9 @@ def get_iam_policy( body = {"options": {"requestedPolicyVersion": 1}} path = "{}:getIamPolicy".format(table.path) - span_creator = SpanCreator() - attributes = {"path": path, "body": body} - with span_creator.create( - name="BigQuery.getIAMPolicy", attributes=attributes, client=self + span_attributes = {"path": path, "body": body} + with create_span( + name="BigQuery.getIAMPolicy", attributes=span_attributes, client=self ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, @@ -685,10 +679,9 @@ def set_iam_policy( body["updateMask"] = updateMask path = "{}:setIamPolicy".format(table.path) - span_creator = SpanCreator() - attributes = {"path": path, "body": body} - with span_creator.create( - name="BigQuery.setIAMPolicy", attributes=attributes, client=self + span_attributes = {"path": path, "body": body} + with create_span( + name="BigQuery.setIAMPolicy", attributes=span_attributes, client=self ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, @@ -705,10 +698,9 @@ def test_iam_permissions( body = {"permissions": permissions} path = "{}:testIamPermissions".format(table.path) - span_creator = SpanCreator() - attributes = {"path": path, "body": body} - with span_creator.create( - name="BigQuery.testIAMPolicy", attributes=attributes, client=self + span_attributes = {"path": path, "body": body} + with create_span( + name="BigQuery.testIamPermissions", attributes=span_attributes, client=self ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, @@ -741,11 +733,10 @@ def get_model(self, model_ref, retry=DEFAULT_RETRY, timeout=None): model_ref = ModelReference.from_string( model_ref, default_project=self.project ) - span_creator = SpanCreator() path = model_ref.path - attributes = {"path": path} - with span_creator.create( - name="BigQuery.getModel", attributes=attributes, client=self + 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 @@ -781,10 +772,9 @@ def get_routine(self, routine_ref, retry=DEFAULT_RETRY, timeout=None): routine_ref, default_project=self.project ) path = routine_ref.path - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.getRoutine", attributes=attributes, client=self + 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 @@ -816,10 +806,9 @@ 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_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.getTable", attributes=attributes, client=self + 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 @@ -861,11 +850,10 @@ def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): else: headers = None path = dataset.path - span_creator = SpanCreator() - attributes = {"path": path, "fields": fields} + span_attributes = {"path": path, "fields": fields} - with span_creator.create( - name="BigQuery.updateDataset", attributes=attributes, client=self + with create_span( + name="BigQuery.updateDataset", attributes=span_attributes, client=self ): api_response = self._call_api( retry, @@ -911,11 +899,10 @@ def update_model(self, model, fields, retry=DEFAULT_RETRY, timeout=None): else: headers = None path = model.path - span_creator = SpanCreator() - attributes = {"path": path, "fields": fields} + span_attributes = {"path": path, "fields": fields} - with span_creator.create( - name="BigQuery.updateModel", attributes=attributes, client=self + with create_span( + name="BigQuery.updateModel", attributes=span_attributes, client=self ): api_response = self._call_api( retry, @@ -972,11 +959,10 @@ def update_routine(self, routine, fields, retry=DEFAULT_RETRY, timeout=None): partial["routineReference"] = routine.reference.to_api_repr() path = routine.path - span_creator = SpanCreator() - attributes = {"path": path, "fields": fields} + span_attributes = {"path": path, "fields": fields} - with span_creator.create( - name="BigQuery.updateRountine", attributes=attributes, client=self + with create_span( + name="BigQuery.updateRoutine", attributes=span_attributes, client=self ): api_response = self._call_api( retry, @@ -1023,11 +1009,10 @@ def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): headers = None path = table.path - span_creator = SpanCreator() - attributes = {"path": path, "fields": fields} + span_attributes = {"path": path, "fields": fields} - with span_creator.create( - name="BigQuery.updateTable", attributes=attributes, client=self + with create_span( + name="BigQuery.updateTable", attributes=span_attributes, client=self ): api_response = self._call_api( retry, @@ -1092,14 +1077,13 @@ def list_models( raise TypeError("dataset must be a Dataset, DatasetReference, or string") path = "%s/models" % dataset.path - span_creator = SpanCreator() - attributes = { + span_attributes = { "path": path, "page_token": page_token, "max_results": max_results, } - with span_creator.create( - name="BigQuery.listModels", attributes=attributes, client=self + with create_span( + name="BigQuery.listModels", attributes=span_attributes, client=self ): api_request = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( @@ -1167,14 +1151,13 @@ def list_routines( raise TypeError("dataset must be a Dataset, DatasetReference, or string") path = "{}/routines".format(dataset.path) - span_creator = SpanCreator() - attributes = { + span_attributes = { "path": path, "page_token": page_token, "max_results": max_results, } - with span_creator.create( - name="BigQuery.listRoutines", attributes=attributes, client=self + with create_span( + name="BigQuery.listRoutines", attributes=span_attributes, client=self ): api_requests = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( @@ -1242,14 +1225,13 @@ def list_tables( raise TypeError("dataset must be a Dataset, DatasetReference, or string") path = "%s/tables" % dataset.path - span_creator = SpanCreator() - attributes = { + span_attributes = { "path": path, "page_token": page_token, "max_results": max_results, } - with span_creator.create( - name="BigQuery.listTables", attributes=attributes, client=self + with create_span( + name="BigQuery.listTables", attributes=span_attributes, client=self ): api_requests = functools.partial(self._call_api, retry, timeout=timeout) @@ -1315,10 +1297,9 @@ def delete_dataset( params["deleteContents"] = "true" try: - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.deleteDataset", attributes=attributes, client=self + span_attributes = {"path": path} + with create_span( + name="BigQuery.deleteDataset", attributes=span_attributes, client=self ): self._call_api( retry, @@ -1366,10 +1347,9 @@ def delete_model( path = model.path try: - span_creator = SpanCreator() - attributes = {"not_found_okay": not_found_ok, "path": path} - with span_creator.create( - name="BigQuery.deleteModel", attributes=attributes, client=self + span_attributes = {"not_found_okay": not_found_ok, "path": path} + with create_span( + name="BigQuery.deleteModel", attributes=span_attributes, client=self ): self._call_api(retry, method="DELETE", path=path, timeout=timeout) except google.api_core.exceptions.NotFound: @@ -1413,10 +1393,9 @@ def delete_routine( raise TypeError("routine must be a Routine or a RoutineReference") try: - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.deleteRoutine", attributes=attributes, client=self + 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) except google.api_core.exceptions.NotFound: @@ -1456,10 +1435,9 @@ def delete_table( try: path = table.path - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.deleteTable", attributes=attributes, client=self + 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) except google.api_core.exceptions.NotFound: @@ -1509,11 +1487,10 @@ def _get_query_results( # This call is typically made in a polling loop that checks whether the # job is complete (from QueryJob.done(), called ultimately from # QueryJob.result()). So we don't need to poll here. - span_creator = SpanCreator() - attributes = {"path": path, "job_id": job_id, "location": location} + span_attributes = {"path": path, "job_id": job_id, "location": location} - with span_creator.create( - name="BigQuery.getQueryResults", attributes=attributes, client=self + with create_span( + name="BigQuery.getQueryResults", attributes=span_attributes, client=self ): resource = self._call_api( retry, @@ -1664,12 +1641,11 @@ def get_job( extra_params["location"] = location path = "/projects/{}/jobs/{}".format(project, job_id) - span_creator = SpanCreator() - attributes = {"path": path, "job_id": job_id, "location": location} + span_attributes = {"path": path, "job_id": job_id, "location": location} - with span_creator.create( - name="BigQuery.getJob", attributes=attributes, client=self + with create_span( + name="BigQuery.getJob", attributes=span_attributes, client=self ): resource = self._call_api( retry, @@ -1723,12 +1699,11 @@ def cancel_job( extra_params["location"] = location path = "/projects/{}/jobs/{}/cancel".format(project, job_id) - span_creator = SpanCreator() - attributes = {"path": path, "job_id": job_id, "location": location} + span_attributes = {"path": path, "job_id": job_id, "location": location} - with span_creator.create( - name="BigQuery.CancelJob", attributes=attributes, client=self + with create_span( + name="BigQuery.cancelJob", attributes=span_attributes, client=self ): resource = self._call_api( retry, @@ -1826,15 +1801,14 @@ def list_jobs( project = self.project path = "/projects/%s/jobs" % (project,) - span_creator = SpanCreator() - attributes = { + span_attributes = { "path": path, "all_users": all_users, "state_filter": state_filter, } - with span_creator.create( - name="BigQuery.listJobs", attributes=attributes, client=self + with create_span( + name="BigQuery.listJobs", attributes=span_attributes, client=self ): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( @@ -2924,10 +2898,9 @@ def insert_rows_json( path = "%s/insertAll" % table.path # We can always retry, because every row has an insert ID. - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.insertRowsJson", attributes=attributes, client=self + 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, @@ -3069,17 +3042,16 @@ def list_rows( params["selectedFields"] = ",".join(field.name for field in selected_fields) if start_index is not None: params["startIndex"] = start_index - span_creator = SpanCreator() path = "%s/data" % (table.path,) - attributes = { + span_attributes = { "path": path, "page_size": page_size, "page_token": page_token, "start_index": start_index, } - with span_creator.create( - name="BigQuery.listRows", attributes=attributes, client=self + with create_span( + name="BigQuery.listRows", attributes=span_attributes, client=self ): api_request = functools.partial(self._call_api, retry, timeout=timeout) row_iterator = RowIterator( diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index 5ee86cb39..6db71b551 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -50,7 +50,7 @@ from google.cloud.bigquery.table import TableReference from google.cloud.bigquery.table import Table from google.cloud.bigquery.table import TimePartitioning -from google.cloud.bigquery.opentelemetry_tracing import SpanCreator +from google.cloud.bigquery.opentelemetry_tracing import create_span _DONE_STATE = "DONE" @@ -636,10 +636,9 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None): # jobs.insert is idempotent because we ensure that every new # job has an ID. - attributes = {"path": path} - span_creator = SpanCreator() - with span_creator.create( - name="BigQuery.job.begin", attributes=attributes, job_ref=self + span_attributes = {"path": path} + with create_span( + name="BigQuery.job.begin", attributes=span_attributes, job_ref=self ): api_response = client._call_api( retry, @@ -676,10 +675,9 @@ def exists(self, client=None, retry=DEFAULT_RETRY, timeout=None): extra_params["location"] = self.location try: - span_creator = SpanCreator() - attributes = {"path": self.path} - with span_creator.create( - name="BigQuery.job.exists", attributes=attributes, job_ref=self + span_attributes = {"path": self.path} + with create_span( + name="BigQuery.job.exists", attributes=span_attributes, job_ref=self ): client._call_api( retry, @@ -714,10 +712,9 @@ def reload(self, client=None, retry=DEFAULT_RETRY, timeout=None): extra_params = {} if self.location: extra_params["location"] = self.location - span_creator = SpanCreator() - attributes = {"path": self.path} - with span_creator.create( - name="BigQuery.job.reload", attributes=attributes, job_ref=self + span_attributes = {"path": self.path} + with create_span( + name="BigQuery.job.reload", attributes=span_attributes, job_ref=self ): api_response = client._call_api( retry, @@ -753,10 +750,9 @@ def cancel(self, client=None, retry=DEFAULT_RETRY, timeout=None): extra_params["location"] = self.location path = "{}/cancel".format(self.path) - span_creator = SpanCreator() - attributes = {"path": path} - with span_creator.create( - name="BigQuery.job.cancel", attributes=attributes, job_ref=self + span_attributes = {"path": path} + with create_span( + name="BigQuery.job.cancel", attributes=span_attributes, job_ref=self ): api_response = client._call_api( retry, diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index e2d8f483f..4d05b7171 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -26,63 +26,62 @@ except ImportError: Logger.info( - "This service instrumented using opentelemetry." - "Opentelemetry could not be imported please" + "This service is instrumented using opentelemetry." + "Opentelemetry could not be imported, please" "add opentelemetry-api and opentelemetry-instrumentation" - "packages in order to get Big Query Tracing data" + "packages, in order to get Big Query Tracing data" ) HAS_OPENTELEMETRY = False -class SpanCreator: - def __init__(self): - # Constructs a span creator with all default attributes - self.opentelemetry_enbled = HAS_OPENTELEMETRY - self.attributes = { - "db.system": "bigquery", - } - - @contextmanager - def create(self, name, attributes=None, client=None, job_ref=None): - if not self.opentelemetry_enbled: - yield None - return - - tracer = trace.get_tracer(__name__) - - if client: - self.set_client_attributes(client) - - elif job_ref: - self.set_job_attributes(job_ref) - - if attributes: - self.attributes.update(attributes) - # yield new span value - with tracer.start_as_current_span( - name=name, attributes=self.attributes - ) as span: - try: - yield span - except GoogleAPICallError as error: - if error.code is not None: - span.set_status(Status(http_status_to_canonical_code(error.code))) - raise - - def set_client_attributes(self, client): - self.attributes["db.name"] = client.project - self.attributes["location"] = client.location - - def set_job_attributes(self, job_ref): - self.attributes["db.name"] = job_ref.project - self.attributes["location"] = job_ref.location - self.attributes["num_child_jobs"] = str(job_ref.num_child_jobs) - self.attributes["job_id"] = job_ref.job_id - self.attributes["parent_job_id"] = job_ref.parent_job_id - self.attributes["timeCreated"] = job_ref.created - self.attributes["timeStarted"] = job_ref.started - self.attributes["timeEnded"] = job_ref.ended - self.attributes["errors"] = job_ref.errors - self.attributes["errorResult"] = job_ref.error_result - self.attributes["state"] = job_ref.state +@contextmanager +def create_span(name, attributes=None, client=None, job_ref=None): + if not HAS_OPENTELEMETRY: + yield None + return + + default_attributes = { + "db.system": "BigQuery", + } + + tracer = trace.get_tracer(__name__) + + if client: + client_attributes = _set_client_attributes(client) + default_attributes.update(client_attributes) + elif job_ref: + job_attributes = _set_job_attributes(job_ref) + default_attributes.update(job_attributes) + + if attributes: + default_attributes.update(attributes) + + # yield new span value + with tracer.start_as_current_span(name=name, attributes=default_attributes) as span: + try: + yield span + except GoogleAPICallError as error: + if error.code is not None: + span.set_status(Status(http_status_to_canonical_code(error.code))) + raise + + +def _set_client_attributes(client): + return {"db.name": client.project, "location": client.location} + + +def _set_job_attributes(job_ref): + return { + "db.name": job_ref.project, + "location": job_ref.location, + "num_child_jobs": str(job_ref.num_child_jobs), + "job_id": job_ref.job_id, + "parent_job_id": job_ref.parent_job_id, + "timeCreated": job_ref.created, + "timeStarted": job_ref.started, + "timeEnded": job_ref.ended, + "errors": job_ref.errors, + "errorResult": job_ref.error_result, + "state": job_ref.state, + } diff --git a/setup.py b/setup.py index f391143d3..9fdf8c9ab 100644 --- a/setup.py +++ b/setup.py @@ -62,6 +62,12 @@ "llvmlite <= 0.33.0;python_version>='3.6'", "llvmlite <= 0.31.0;python_version<'3.6'", ], + "opentelemetry": [ + "opentelemetry-api", + "opentelemetry-sdk", + "opentelemetry-instrumentaion", + "opentelemetry-exporter-cloud-trace", + ], } all_extras = [] diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index eb7630eb4..1dd99013a 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -29,17 +29,17 @@ def test_opentelemetry_not_installed(setup): temp_module = sys.modules["opentelemetry"] sys.modules["opentelemetry"] = None reload(opentelemetry_tracing) - span_creator = opentelemetry_tracing.SpanCreator() - with span_creator.create("No-op for opentelemetry") as span: + with opentelemetry_tracing.create_span("No-op for opentelemetry") as span: assert span is None sys.modules["opentelemetry"] = temp_module reload(opentelemetry_tracing) def test_opentelemetry_success(setup): - span_creator = opentelemetry_tracing.SpanCreator() - expected_attributes = {"foo": "bar", "db.system": "bigquery"} - with span_creator.create(TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES,) as span: + expected_attributes = {"foo": "bar", "db.system": "BigQuery"} + with opentelemetry_tracing.create_span( + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, + ) as span: if span is None: span_list = setup.get_finished_spans() print(span_list) @@ -56,14 +56,13 @@ def test_default_client_attributes(setup): project="test_project", credentials=mock_credentials, location="test_location" ) - span_creator = opentelemetry_tracing.SpanCreator() expected_attributes = { "foo": "bar", - "db.system": "bigquery", + "db.system": "BigQuery", "db.name": "test_project", "location": "test_location", } - with span_creator.create( + with opentelemetry_tracing.create_span( TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client ) as span: if span is None: @@ -88,10 +87,8 @@ def test_default_job_attributes(setup): ) test_job = job._AsyncJob(job_id=test_job_reference, client=test_client) - span_creator = opentelemetry_tracing.SpanCreator() - # TODO add assertions for attributes being equal expected_attributes = { - "db.system": "bigquery", + "db.system": "BigQuery", "db.name": "test_project_id", "location": "test_location", "num_child_jobs": "0", @@ -99,7 +96,7 @@ def test_default_job_attributes(setup): "foo": "bar", } - with span_creator.create( + with opentelemetry_tracing.create_span( TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job ) as span: if span is None: From 4dbb6923d26a239f1d1147856c506c9ac6b29536 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Tue, 11 Aug 2020 21:26:39 -0400 Subject: [PATCH 12/41] adding suggested changes --- README.rst | 4 +- .../cloud/bigquery/opentelemetry_tracing.py | 43 +++++++++++++++---- noxfile.py | 9 +--- setup.py | 8 ++-- 4 files changed, 42 insertions(+), 22 deletions(-) diff --git a/README.rst b/README.rst index 145b8f1c8..3dab8c68a 100644 --- a/README.rst +++ b/README.rst @@ -134,10 +134,10 @@ can be found here: ) In this example all tracing data will be published to the Google -Cloud trace console. For more information on OpenTelemetry, please consult the `OpenTelemetry documentation`_. +`Cloud Trace`_ console. For more information on OpenTelemetry, please consult the `OpenTelemetry documentation`_. .. _OpenTelemetry documentation: https://opentelemetry-python.readthedocs.io - +.. _Cloud Trace: https://cloud.google.com/trace diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 4d05b7171..99bb478b1 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -35,32 +35,59 @@ HAS_OPENTELEMETRY = False +__default_attributes = { + "db.system": "BigQuery", +} + + @contextmanager def create_span(name, attributes=None, client=None, job_ref=None): + """Creates a ContextManager for a Span to be exported to the configured exporter. If no configuration + exists yields None. + + Args: + name (str): Name that will be set for the span being created + attributes(Optional[dict]): + Additional attributes that pertain to + the specific API call (i.e. not a default attribute) + client (Optional[google.cloud.bigquery.client.Client]): + Pass in a Client object to extract any attributes that may be + relevant to it and add them to the created spans. + job_ref(Optional[google.cloud.bigquery.job._AsyncJob]) + Pass in a _AsyncJob object to extract any attributes that may be + relevant to it and add them to the created spans. + + Yields: + opentelemetry.trace.Span: Yields the newly created Span. + + Raises: + google.api_core.exceptions.GoogleAPICallError: + Raised if a span could not be yielded or issue with call to + OpenTelemetry. + """ if not HAS_OPENTELEMETRY: yield None return - default_attributes = { - "db.system": "BigQuery", - } - tracer = trace.get_tracer(__name__) if client: client_attributes = _set_client_attributes(client) - default_attributes.update(client_attributes) + __default_attributes.update(client_attributes) elif job_ref: job_attributes = _set_job_attributes(job_ref) - default_attributes.update(job_attributes) + __default_attributes.update(job_attributes) if attributes: - default_attributes.update(attributes) + __default_attributes.update(attributes) # yield new span value - with tracer.start_as_current_span(name=name, attributes=default_attributes) as span: + with tracer.start_as_current_span( + name=name, attributes=__default_attributes + ) as span: try: yield span + span.set_status(Status(http_status_to_canonical_code(200))) except GoogleAPICallError as error: if error.code is not None: span.set_status(Status(http_status_to_canonical_code(error.code))) diff --git a/noxfile.py b/noxfile.py index 1a7f9c13d..dd0840189 100644 --- a/noxfile.py +++ b/noxfile.py @@ -33,14 +33,7 @@ def default(session): """ # Install all test dependencies, then install local packages in-place. session.install( - "mock", - "pytest", - "google-cloud-testutils", - "pytest-cov", - "freezegun", - "opentelemetry-api", - "opentelemetry-sdk", - "opentelemetry-instrumentation", + "mock", "pytest", "google-cloud-testutils", "pytest-cov", "freezegun", ) session.install("grpcio") diff --git a/setup.py b/setup.py index 833e0d125..3fb0ba7f7 100644 --- a/setup.py +++ b/setup.py @@ -61,10 +61,10 @@ "llvmlite <= 0.31.0;python_version<'3.6'", ], "opentelemetry": [ - "opentelemetry-api", - "opentelemetry-sdk", - "opentelemetry-instrumentaion", - "opentelemetry-exporter-cloud-trace", + "opentelemetry-api == 0.9b0", + "opentelemetry-sdk == 0.9b0", + "opentelemetry-instrumentation == 0.9b0 ", + "opentelemetry-exporter-cloud-trace == 0.110", ], } From 406f61a60c70119d6146d5ced8e18a17a6deff47 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Tue, 11 Aug 2020 21:56:02 -0400 Subject: [PATCH 13/41] removing print statements --- tests/unit/test_opentelemetry_tracing.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 1dd99013a..bd19feab5 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -42,7 +42,6 @@ def test_opentelemetry_success(setup): ) as span: if span is None: span_list = setup.get_finished_spans() - print(span_list) assert len(span_list) == 1 assert span.name == TEST_SPAN_NAME assert span.attributes == expected_attributes @@ -67,7 +66,6 @@ def test_default_client_attributes(setup): ) as span: if span is None: span_list = setup.get_finished_spans() - print(span_list) assert len(span_list) == 1 assert span.name == TEST_SPAN_NAME assert span.attributes == expected_attributes @@ -101,7 +99,6 @@ def test_default_job_attributes(setup): ) as span: if span is None: span_list = setup.get_finished_spans() - print(span_list[0]) assert len(span_list) == 1 assert span.name == TEST_SPAN_NAME assert span.attributes == expected_attributes From c2b9addb5a7a533440228a22b12e89ef173cd2e8 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Wed, 12 Aug 2020 14:59:56 -0400 Subject: [PATCH 14/41] setting same version across all OT [ackages and other reccommended changes --- tests/unit/test_opentelemetry_tracing.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index bd19feab5..18da91536 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -12,7 +12,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter TEST_SPAN_NAME = "bar" -TEST_SPAN_ATTRIBUTES = {"foo": "bar"} +TEST_SPAN_ATTRIBUTES = {"foo": "baz"} @pytest.fixture @@ -36,7 +36,7 @@ def test_opentelemetry_not_installed(setup): def test_opentelemetry_success(setup): - expected_attributes = {"foo": "bar", "db.system": "BigQuery"} + expected_attributes = {"foo": "baz", "db.system": "BigQuery"} with opentelemetry_tracing.create_span( TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, ) as span: @@ -56,7 +56,7 @@ def test_default_client_attributes(setup): ) expected_attributes = { - "foo": "bar", + "foo": "baz", "db.system": "BigQuery", "db.name": "test_project", "location": "test_location", @@ -91,7 +91,7 @@ def test_default_job_attributes(setup): "location": "test_location", "num_child_jobs": "0", "job_id": "test_job_id", - "foo": "bar", + "foo": "baz", } with opentelemetry_tracing.create_span( From d4259099c5d241b6d800107ca47a85363c4f2dce Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Wed, 12 Aug 2020 15:00:33 -0400 Subject: [PATCH 15/41] suggested changes --- README.rst | 3 +-- google/cloud/bigquery/client.py | 8 ++++---- google/cloud/bigquery/opentelemetry_tracing.py | 10 +++++----- setup.py | 2 +- 4 files changed, 11 insertions(+), 12 deletions(-) diff --git a/README.rst b/README.rst index 3dab8c68a..dd74f0b69 100644 --- a/README.rst +++ b/README.rst @@ -106,8 +106,7 @@ Perform a query Instrumenting With OpenTelemetry -------------------------------- -This application uses `OpenTelemetry`_ -to output tracing data from +This application uses `OpenTelemetry`_ to output tracing data from API calls to BigQuery. To enable OpenTelemetry tracing in the BigQuery client the following PyPI packages need to be installed: diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index ad28cb87e..f50ae93b0 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -288,7 +288,7 @@ def list_projects( path = "/projects" span_attributes = {"path": path} with create_span( - name="BigQuery.list_projects", attributes=span_attributes, client=self + name="BigQuery.listProjects", attributes=span_attributes, client=self ): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( @@ -358,7 +358,7 @@ def list_datasets( span_attributes = {"path": path, "page_token": page_token} with create_span( - name="BigQuery.list_datasets", attributes=span_attributes, client=self + name="BigQuery.listDatasets", attributes=span_attributes, client=self ): api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( @@ -635,7 +635,7 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None): path = dataset_ref.path span_attributes = {"path": path} with create_span( - name="BigQuery.createTable", attributes=span_attributes, client=self + name="BigQuery.getDataset", attributes=span_attributes, client=self ): api_response = self._call_api( retry, method="GET", path=path, timeout=timeout @@ -700,7 +700,7 @@ def test_iam_permissions( path = "{}:testIamPermissions".format(table.path) span_attributes = {"path": path, "body": body} with create_span( - name="BigQuery.testIamPermissions", attributes=span_attributes, client=self + name="BigQuery.testIAMPermissions", attributes=span_attributes, client=self ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 99bb478b1..a656f0a63 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -35,7 +35,7 @@ HAS_OPENTELEMETRY = False -__default_attributes = { +_default_attributes = { "db.system": "BigQuery", } @@ -73,17 +73,17 @@ def create_span(name, attributes=None, client=None, job_ref=None): if client: client_attributes = _set_client_attributes(client) - __default_attributes.update(client_attributes) + _default_attributes.update(client_attributes) elif job_ref: job_attributes = _set_job_attributes(job_ref) - __default_attributes.update(job_attributes) + _default_attributes.update(job_attributes) if attributes: - __default_attributes.update(attributes) + _default_attributes.update(attributes) # yield new span value with tracer.start_as_current_span( - name=name, attributes=__default_attributes + name=name, attributes=_default_attributes ) as span: try: yield span diff --git a/setup.py b/setup.py index 3fb0ba7f7..c320b2c4e 100644 --- a/setup.py +++ b/setup.py @@ -64,7 +64,7 @@ "opentelemetry-api == 0.9b0", "opentelemetry-sdk == 0.9b0", "opentelemetry-instrumentation == 0.9b0 ", - "opentelemetry-exporter-cloud-trace == 0.110", + "opentelemetry-exporter-cloud-trace == 0.9b0", ], } From 64ec5fdf80a1f26e76fc77f9fe802d26b414ae43 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Thu, 13 Aug 2020 08:51:40 -0400 Subject: [PATCH 16/41] fixing packages issue in nox and updating documentation --- README.rst | 2 +- noxfile.py | 8 ++++++++ setup.py | 6 ------ 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/README.rst b/README.rst index dd74f0b69..4cfde7076 100644 --- a/README.rst +++ b/README.rst @@ -114,7 +114,7 @@ the BigQuery client the following PyPI packages need to be installed: .. code-block:: console - pip install google-cloud-bigquery[opentelemetry] + pip install opentelemetry-api opentelemetry-sdk opentelemetry-instrumentation opentelemetry-exporter-google-cloud After installation, OpenTelemetry can be used in the BigQuery client and in BigQuery jobs. First, however, an exporter must be diff --git a/noxfile.py b/noxfile.py index dd0840189..fa16d9770 100644 --- a/noxfile.py +++ b/noxfile.py @@ -48,6 +48,14 @@ def default(session): else: session.install("ipython") + # opentelemetry was not added to [all] because opentelemetry does not support Python 2 + # exporter does not need to be in nox thus it has been added to README documentation + if session.python == "2.7": + session.install( + "opentelemetry-api == 0.9b0", + "opentelemetry-sdk == 0.9b0", + "opentelemetry-instrumentation == 0.9b0 ", + ) # Run py.test against the unit tests. session.run( "py.test", diff --git a/setup.py b/setup.py index c320b2c4e..b00b2cbe5 100644 --- a/setup.py +++ b/setup.py @@ -60,12 +60,6 @@ "llvmlite <= 0.33.0;python_version>='3.6'", "llvmlite <= 0.31.0;python_version<'3.6'", ], - "opentelemetry": [ - "opentelemetry-api == 0.9b0", - "opentelemetry-sdk == 0.9b0", - "opentelemetry-instrumentation == 0.9b0 ", - "opentelemetry-exporter-cloud-trace == 0.9b0", - ], } all_extras = [] From 73230dc0c0a2b4cabf10725c85492474ec507b7a Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Thu, 13 Aug 2020 13:55:15 -0400 Subject: [PATCH 17/41] fixing module install issue --- noxfile.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/noxfile.py b/noxfile.py index fa16d9770..c4c0ceb7b 100644 --- a/noxfile.py +++ b/noxfile.py @@ -50,11 +50,11 @@ def default(session): # opentelemetry was not added to [all] because opentelemetry does not support Python 2 # exporter does not need to be in nox thus it has been added to README documentation - if session.python == "2.7": + if session.python != "2.7": session.install( - "opentelemetry-api == 0.9b0", - "opentelemetry-sdk == 0.9b0", - "opentelemetry-instrumentation == 0.9b0 ", + "opentelemetry-api==0.9b0", + "opentelemetry-sdk==0.9b0", + "opentelemetry-instrumentation==0.9b0 ", ) # Run py.test against the unit tests. session.run( From ae04fc6fe10a4f6bb8b8d84835e56bdde0883a1c Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Fri, 14 Aug 2020 00:00:08 -0400 Subject: [PATCH 18/41] restructuring design for testing adding first layer of tests (some still failing) --- google/cloud/bigquery/client.py | 2 +- .../cloud/bigquery/opentelemetry_tracing.py | 41 +- tests/unit/test_client.py | 564 +++++++++++++----- 3 files changed, 435 insertions(+), 172 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index f50ae93b0..d052e60a5 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -679,7 +679,7 @@ def set_iam_policy( body["updateMask"] = updateMask path = "{}:setIamPolicy".format(table.path) - span_attributes = {"path": path, "body": body} + span_attributes = {"path": path} with create_span( name="BigQuery.setIAMPolicy", attributes=span_attributes, client=self ): diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index a656f0a63..8fa8ceb65 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -39,9 +39,31 @@ "db.system": "BigQuery", } +tracer = trace.get_tracer(__name__) + + @contextmanager def create_span(name, attributes=None, client=None, job_ref=None): + # yield new span value + if not HAS_OPENTELEMETRY: + yield None + return + + final_attributes = _get_final_span_attributes(attributes, client, job_ref) + with tracer.start_as_current_span( + name=name, attributes=final_attributes + ) as span: + try: + yield span + span.set_status(Status(http_status_to_canonical_code(200))) + except GoogleAPICallError as error: + if error.code is not None: + span.set_status(Status(http_status_to_canonical_code(error.code))) + raise + + +def _get_final_span_attributes(attributes=None, client=None, job_ref=None): """Creates a ContextManager for a Span to be exported to the configured exporter. If no configuration exists yields None. @@ -65,11 +87,7 @@ def create_span(name, attributes=None, client=None, job_ref=None): Raised if a span could not be yielded or issue with call to OpenTelemetry. """ - if not HAS_OPENTELEMETRY: - yield None - return - tracer = trace.get_tracer(__name__) if client: client_attributes = _set_client_attributes(client) @@ -81,17 +99,7 @@ def create_span(name, attributes=None, client=None, job_ref=None): if attributes: _default_attributes.update(attributes) - # yield new span value - with tracer.start_as_current_span( - name=name, attributes=_default_attributes - ) as span: - try: - yield span - span.set_status(Status(http_status_to_canonical_code(200))) - except GoogleAPICallError as error: - if error.code is not None: - span.set_status(Status(http_status_to_canonical_code(error.code))) - raise + return _default_attributes def _set_client_attributes(client): @@ -112,3 +120,6 @@ def _set_job_attributes(job_ref): "errorResult": job_ref.error_result, "state": job_ref.state, } + + + diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 8b63f7e57..a4ba43e76 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -57,6 +57,7 @@ bigquery_storage_v1 = None from test_utils.imports import maybe_fail_import from tests.unit.helpers import make_connection +import pdb PANDAS_MINIUM_VERSION = pkg_resources.parse_version("1.0.0") PANDAS_INSTALLED_VERSION = pkg_resources.get_distribution("pandas").parsed_version @@ -246,20 +247,23 @@ def test__get_query_results_miss_w_explicit_project_and_timeout(self): creds = _make_credentials() client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection() - + path = "/projects/other-project/queries/nothere" with self.assertRaises(NotFound): - client._get_query_results( - "nothere", - None, - project="other-project", - location=self.LOCATION, - timeout_ms=500, - timeout=42, - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client._get_query_results( + "nothere", + None, + project="other-project", + location=self.LOCATION, + timeout_ms=500, + timeout=42, + ) + final_attributes.assert_called_once_with({'path': path}, client, None) + conn.api_request.assert_called_once_with( method="GET", - path="/projects/other-project/queries/nothere", + path=path, query_params={"maxResults": 0, "timeoutMs": 500, "location": self.LOCATION}, timeout=42, ) @@ -315,9 +319,9 @@ def test_get_service_account_email(self): email = "bq-123@bigquery-encryption.iam.gserviceaccount.com" resource = {"kind": "bigquery#getServiceAccountResponse", "email": email} conn = client._connection = make_connection(resource) - - service_account_email = client.get_service_account_email(timeout=7.5) - + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + service_account_email = client.get_service_account_email(timeout=7.5) + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_once_with(method="GET", path=path, timeout=7.5) self.assertEqual(service_account_email, email) @@ -330,9 +334,9 @@ def test_get_service_account_email_w_alternate_project(self): email = "bq-123@bigquery-encryption.iam.gserviceaccount.com" resource = {"kind": "bigquery#getServiceAccountResponse", "email": email} conn = client._connection = make_connection(resource) - - service_account_email = client.get_service_account_email(project=project) - + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + service_account_email = client.get_service_account_email(project=project) + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_once_with(method="GET", path=path, timeout=None) self.assertEqual(service_account_email, email) @@ -357,10 +361,11 @@ def test_get_service_account_email_w_custom_retry(self): ) with api_request_patcher as fake_api_request: - service_account_email = client.get_service_account_email( - retry=retry, timeout=7.5 - ) - + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + service_account_email = client.get_service_account_email( + retry=retry, timeout=7.5 + ) + final_attributes.assert_called_once_with({'path': api_path}, client, None) self.assertEqual( service_account_email, "bq-123@bigquery-encryption.iam.gserviceaccount.com" ) @@ -400,8 +405,9 @@ 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() + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_projects() + final_attributes.assert_called_once_with({'path': '/projects'}, client, None) page = six.next(iterator.pages) projects = list(page) token = iterator.next_page_token @@ -428,8 +434,10 @@ def test_list_projects_w_timeout(self): creds = _make_credentials() client = self._make_one(PROJECT_1, creds) conn = client._connection = make_connection(DATA) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_projects(timeout=7.5) + final_attributes.assert_called_once_with({'path': '/projects'}, client, None) - iterator = client.list_projects(timeout=7.5) six.next(iterator.pages) conn.api_request.assert_called_once_with( @@ -443,7 +451,10 @@ def test_list_projects_explicit_response_missing_projects_key(self): client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - iterator = client.list_projects(max_results=3, page_token=TOKEN) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_projects(max_results=3, page_token=TOKEN) + final_attributes.assert_called_once_with({'path': '/projects'}, client, None) + page = six.next(iterator.pages) projects = list(page) token = iterator.next_page_token @@ -491,8 +502,10 @@ def test_list_datasets_defaults(self): creds = _make_credentials() client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_datasets() + final_attributes.assert_called_once_with({'path': PATH, 'page_token':TOKEN}, client, None) - iterator = client.list_datasets() page = six.next(iterator.pages) datasets = list(page) token = iterator.next_page_token @@ -513,7 +526,10 @@ 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() conn.api_request.assert_called_once_with( method="GET", @@ -531,9 +547,13 @@ def test_list_datasets_explicit_response_missing_datasets_key(self): client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - iterator = client.list_datasets( - include_all=True, filter=FILTER, max_results=3, page_token=TOKEN - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_datasets( + include_all=True, filter=FILTER, max_results=3, page_token=TOKEN + ) + final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': TOKEN}, client, None) + + page = six.next(iterator.pages) datasets = list(page) token = iterator.next_page_token @@ -612,8 +632,11 @@ def test_get_dataset(self): } conn = client._connection = make_connection(resource) dataset_ref = DatasetReference(self.PROJECT, self.DS_ID) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + dataset = client.get_dataset(dataset_ref, timeout=7.5) + + final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) - dataset = client.get_dataset(dataset_ref, timeout=7.5) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 @@ -625,35 +648,50 @@ def test_get_dataset(self): # Not a cloud API exception (missing 'errors' field). client._connection = make_connection(Exception(""), resource) with self.assertRaises(Exception): - client.get_dataset(dataset_ref) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.get_dataset(dataset_ref) + final_attributes.assert_called_once_with({'path': path}, client, None) + # Zero-length errors field. client._connection = make_connection(ServerError(""), resource) with self.assertRaises(ServerError): - client.get_dataset(dataset_ref) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.get_dataset(dataset_ref) + final_attributes.assert_called_once_with({'path': path}, client, None) + # Non-retryable reason. client._connection = make_connection( ServerError("", errors=[{"reason": "serious"}]), resource ) with self.assertRaises(ServerError): - client.get_dataset(dataset_ref) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.get_dataset(dataset_ref) + + final_attributes.assert_called_once_with({'path': path}, client, None) # Retryable reason, but retry is disabled. client._connection = make_connection( ServerError("", errors=[{"reason": "backendError"}]), resource ) with self.assertRaises(ServerError): - client.get_dataset(dataset_ref, retry=None) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.get_dataset(dataset_ref, retry=None) + + final_attributes.assert_called_once_with({'path': path}, client, None) # Retryable reason, default retry: success. client._connection = make_connection( ServerError("", errors=[{"reason": "backendError"}]), resource ) - dataset = client.get_dataset( - # Test with a string for dataset ID. - dataset_ref.dataset_id - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + dataset = client.get_dataset( + # Test with a string for dataset ID. + dataset_ref.dataset_id + ) + final_attributes.assert_called_once_with({'path': path}, client, None) + self.assertEqual(dataset.dataset_id, self.DS_ID) @unittest.skipIf( @@ -713,8 +751,11 @@ def test_create_dataset_minimal(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) before = Dataset(ds_ref) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + after = client.create_dataset(before, timeout=7.5) + + final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) - after = client.create_dataset(before, timeout=7.5) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -775,8 +816,9 @@ def test_create_dataset_w_attrs(self): before.default_table_expiration_ms = 3600 before.location = LOCATION before.labels = LABELS - - after = client.create_dataset(before) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + after = client.create_dataset(before) + final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -826,8 +868,10 @@ def test_create_dataset_w_custom_property(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) before = Dataset(ds_ref) before._properties["newAlphaProperty"] = "unreleased property" + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + after = client.create_dataset(before) - after = client.create_dataset(before) + final_attributes.assert_called_once_with({'path': path}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -865,8 +909,10 @@ def test_create_dataset_w_client_location_wo_dataset_location(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) before = Dataset(ds_ref) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + after = client.create_dataset(before) - after = client.create_dataset(before) + final_attributes.assert_called_once_with({'path': PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -908,8 +954,9 @@ def test_create_dataset_w_client_location_w_dataset_location(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) before = Dataset(ds_ref) before.location = OTHER_LOCATION - - after = client.create_dataset(before) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + after = client.create_dataset(before) + final_attributes.assert_called_once_with({'path': PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -944,8 +991,10 @@ def test_create_dataset_w_reference(self): project=self.PROJECT, credentials=creds, location=self.LOCATION ) conn = client._connection = make_connection(resource) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + dataset = client.create_dataset(DatasetReference(self.PROJECT, self.DS_ID)) - dataset = client.create_dataset(DatasetReference(self.PROJECT, self.DS_ID)) + final_attributes.assert_called_once_with({'path': path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -980,8 +1029,10 @@ def test_create_dataset_w_fully_qualified_string(self): project=self.PROJECT, credentials=creds, location=self.LOCATION ) conn = client._connection = make_connection(resource) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + dataset = client.create_dataset("{}.{}".format(self.PROJECT, self.DS_ID)) - dataset = client.create_dataset("{}.{}".format(self.PROJECT, self.DS_ID)) + final_attributes.assert_called_once_with({'path': path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -1016,8 +1067,10 @@ def test_create_dataset_w_string(self): project=self.PROJECT, credentials=creds, location=self.LOCATION ) conn = client._connection = make_connection(resource) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + dataset = client.create_dataset(self.DS_ID) - dataset = client.create_dataset(self.DS_ID) + final_attributes.assert_called_once_with({'path': path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -1049,7 +1102,11 @@ def test_create_dataset_alreadyexists_w_exists_ok_false(self): ) with pytest.raises(google.api_core.exceptions.AlreadyExists): - client.create_dataset(self.DS_ID) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.create_dataset(self.DS_ID) + + final_attributes.assert_called_once_with({'path': 'path'}, client, None) + def test_create_dataset_alreadyexists_w_exists_ok_true(self): post_path = "/projects/{}/datasets".format(self.PROJECT) @@ -1067,8 +1124,9 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(self): conn = client._connection = make_connection( google.api_core.exceptions.AlreadyExists("dataset already exists"), resource ) - - dataset = client.create_dataset(self.DS_ID, exists_ok=True) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + dataset = client.create_dataset(self.DS_ID, exists_ok=True) + final_attributes.assert_called_once_with({'path': post_path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -1100,6 +1158,7 @@ def test_create_routine_w_minimal_resource(self): from google.cloud.bigquery.routine import RoutineReference creds = _make_credentials() + path = "/projects/test-routine-project/datasets/test_routines/routines" resource = { "routineReference": { "projectId": "test-routine-project", @@ -1111,12 +1170,15 @@ def test_create_routine_w_minimal_resource(self): conn = client._connection = make_connection(resource) full_routine_id = "test-routine-project.test_routines.minimal_routine" routine = Routine(full_routine_id) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + actual_routine = client.create_routine(routine, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path, 'exists_okay': False}, client, None) - actual_routine = client.create_routine(routine, timeout=7.5) conn.api_request.assert_called_once_with( method="POST", - path="/projects/test-routine-project/datasets/test_routines/routines", + path='path', data=resource, timeout=7.5, ) @@ -1132,11 +1194,15 @@ def test_create_routine_w_conflict(self): conn = client._connection = make_connection( google.api_core.exceptions.AlreadyExists("routine already exists") ) + path = '/projects/test-routine-project/datasets/test_routines/routines' full_routine_id = "test-routine-project.test_routines.minimal_routine" routine = Routine(full_routine_id) with pytest.raises(google.api_core.exceptions.AlreadyExists): - client.create_routine(routine) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.create_routine(routine) + + final_attributes.assert_called_once_with({'path': path, 'exists_okay': False}, client, None) resource = { "routineReference": { @@ -1147,7 +1213,7 @@ def test_create_routine_w_conflict(self): } conn.api_request.assert_called_once_with( method="POST", - path="/projects/test-routine-project/datasets/test_routines/routines", + path=path, data=resource, timeout=None, ) @@ -1164,13 +1230,18 @@ def test_create_routine_w_conflict_exists_ok(self): "routineId": "minimal_routine", } } + path = "/projects/test-routine-project/datasets/test_routines/routines", + conn = client._connection = make_connection( google.api_core.exceptions.AlreadyExists("routine already exists"), resource ) full_routine_id = "test-routine-project.test_routines.minimal_routine" routine = Routine(full_routine_id) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + actual_routine = client.create_routine(routine, exists_ok=True) + + final_attributes.assert_called_once_with({'path': path, 'exists_okay': True}, client, None) - actual_routine = client.create_routine(routine, exists_ok=True) self.assertEqual(actual_routine.project, "test-routine-project") self.assertEqual(actual_routine.dataset_id, "test_routines") @@ -1179,7 +1250,7 @@ def test_create_routine_w_conflict_exists_ok(self): [ mock.call( method="POST", - path="/projects/test-routine-project/datasets/test_routines/routines", + path=path, data=resource, timeout=None, ), @@ -1202,8 +1273,10 @@ def test_create_table_w_day_partition(self): conn = client._connection = make_connection(resource) table = Table(self.TABLE_REF) table.time_partitioning = TimePartitioning() + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table(table, timeout=7.5) - got = client.create_table(table, timeout=7.5) + final_attributes.assert_called_once_with({'path': path, 'dataset_id' : table.dataset_id, 'exists_okay': True}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1235,8 +1308,9 @@ def test_create_table_w_custom_property(self): conn = client._connection = make_connection(resource) table = Table(self.TABLE_REF) table._properties["newAlphaProperty"] = "unreleased property" - - got = client.create_table(table) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table(table) + final_attributes.assert_called_once_with({'path': path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1270,8 +1344,10 @@ def test_create_table_w_encryption_configuration(self): table.encryption_configuration = EncryptionConfiguration( kms_key_name=self.KMS_KEY_NAME ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table(table) - got = client.create_table(table) + final_attributes.assert_called_once_with({'path' : path, 'dataset_id' : table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1300,8 +1376,10 @@ def test_create_table_w_day_partition_and_expire(self): conn = client._connection = make_connection(resource) table = Table(self.TABLE_REF) table.time_partitioning = TimePartitioning(expiration_ms=100) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table(table) - got = client.create_table(table) + final_attributes.assert_called_once_with({'path': path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1359,7 +1437,11 @@ def test_create_table_w_schema_and_query(self): table = Table(self.TABLE_REF, schema=schema) table.view_query = query - got = client.create_table(table) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table(table) + + final_attributes.assert_called_once_with({'path': path, 'dataset_id': table.dataset_id}, client, None) + conn.api_request.assert_called_once_with( method="POST", @@ -1420,7 +1502,10 @@ def test_create_table_w_external(self): ec.autodetect = True table.external_data_configuration = ec - got = client.create_table(table) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table(table) + + final_attributes.assert_called_once_with({'path': path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1454,7 +1539,10 @@ def test_create_table_w_reference(self): resource = self._make_table_resource() conn = client._connection = make_connection(resource) - got = client.create_table(self.TABLE_REF) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table(self.TABLE_REF) + + final_attributes.assert_called_once_with({'path': path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1477,10 +1565,11 @@ def test_create_table_w_fully_qualified_string(self): client = self._make_one(project=self.PROJECT, credentials=creds) resource = self._make_table_resource() conn = client._connection = make_connection(resource) - - got = client.create_table( - "{}.{}.{}".format(self.PROJECT, self.DS_ID, self.TABLE_ID) - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table( + "{}.{}.{}".format(self.PROJECT, self.DS_ID, self.TABLE_ID) + ) + final_attributes.assert_called_once_with({'path': path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1503,8 +1592,10 @@ def test_create_table_w_string(self): client = self._make_one(project=self.PROJECT, credentials=creds) resource = self._make_table_resource() conn = client._connection = make_connection(resource) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - got = client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) + final_attributes.assert_called_once_with({'path': path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1532,7 +1623,11 @@ def test_create_table_alreadyexists_w_exists_ok_false(self): ) with pytest.raises(google.api_core.exceptions.AlreadyExists): - client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) + + final_attributes.assert_called_with({'path': post_path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) + conn.api_request.assert_called_once_with( method="POST", @@ -1562,9 +1657,13 @@ def test_create_table_alreadyexists_w_exists_ok_true(self): google.api_core.exceptions.AlreadyExists("table already exists"), resource ) - got = client.create_table( - "{}.{}".format(self.DS_ID, self.TABLE_ID), exists_ok=True - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.create_table( + "{}.{}".format(self.DS_ID, self.TABLE_ID), exists_ok=True + ) + final_attributes.assert_called_once_with({'path': post_path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) + + self.assertEqual(got.project, self.PROJECT) self.assertEqual(got.dataset_id, self.DS_ID) @@ -1619,7 +1718,11 @@ def test_get_model(self): conn = client._connection = make_connection(resource) model_ref = DatasetReference(self.PROJECT, self.DS_ID).model(self.MODEL_ID) - got = client.get_model(model_ref, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.get_model(model_ref, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path}, client, None) + conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 @@ -1645,7 +1748,10 @@ def test_get_model_w_string(self): conn = client._connection = make_connection(resource) model_id = "{}.{}.{}".format(self.PROJECT, self.DS_ID, self.MODEL_ID) - got = client.get_model(model_id) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + got = client.get_model(model_id) + + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=None @@ -1673,14 +1779,20 @@ def test_get_routine(self): }, "routineType": "SCALAR_FUNCTION", } + path = "/projects/test-routine-project/datasets/test_routines/routines/minimal_routine", + client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(resource) - actual_routine = client.get_routine(routine, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + actual_routine = client.get_routine(routine, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path}, client, None) + conn.api_request.assert_called_once_with( method="GET", - path="/projects/test-routine-project/datasets/test_routines/routines/minimal_routine", + path=path, timeout=7.5, ) self.assertEqual( @@ -1710,7 +1822,10 @@ def test_get_table(self): client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) resource = self._make_table_resource() conn = client._connection = make_connection(resource) - table = client.get_table(self.TABLE_REF, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + table = client.get_table(self.TABLE_REF, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 @@ -1734,8 +1849,6 @@ def test_get_table_sets_user_agent(self): _http=http, ) - client.get_table(self.TABLE_REF) - expected_user_agent = user_agent_override.to_user_agent() http.request.assert_called_once_with( url=mock.ANY, @@ -1786,8 +1899,10 @@ def test_get_iam_policy(self): http = object() client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(RETURNED) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + policy = client.get_iam_policy(self.TABLE_REF, timeout=7.5) - policy = client.get_iam_policy(self.TABLE_REF, timeout=7.5) + final_attributes.assert_called_once_with({'path': PATH}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -1810,6 +1925,7 @@ def test_get_iam_policy_w_invalid_table(self): with self.assertRaises(TypeError): client.get_iam_policy(table_resource_string) + def test_get_iam_policy_w_invalid_version(self): creds = _make_credentials() http = object() @@ -1856,9 +1972,11 @@ def test_set_iam_policy(self): client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(RETURNED) - returned_policy = client.set_iam_policy( - self.TABLE_REF, policy, updateMask=MASK, timeout=7.5 - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + returned_policy = client.set_iam_policy( + self.TABLE_REF, policy, updateMask=MASK, timeout=7.5 + ) + final_attributes.assert_called_once_with({'path': PATH}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -1884,8 +2002,10 @@ def test_set_iam_policy_no_mask(self): http = object() client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(RETURNED) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.set_iam_policy(self.TABLE_REF, policy, timeout=7.5) - client.set_iam_policy(self.TABLE_REF, policy, timeout=7.5) + final_attributes.assert_called_once_with({'path': PATH}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -1902,7 +2022,11 @@ def test_set_iam_policy_invalid_policy(self): client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) with self.assertRaises(TypeError): - client.set_iam_policy(self.TABLE_REF, invalid_policy_repr) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.set_iam_policy(self.TABLE_REF, invalid_policy_repr) + + final_attributes.assert_called_once_with({'path': "{}:setIamPolicy".format(self.TABLE_REF.path)}, client, None) + def test_set_iam_policy_w_invalid_table(self): from google.api_core.iam import Policy @@ -1920,7 +2044,11 @@ def test_set_iam_policy_w_invalid_table(self): ) with self.assertRaises(TypeError): - client.set_iam_policy(table_resource_string, policy) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.set_iam_policy(table_resource_string, policy) + + final_attributes.assert_called_once_with({'path': table_resource_string}, client, None) + def test_test_iam_permissions(self): PATH = "/projects/%s/datasets/%s/tables/%s:testIamPermissions" % ( @@ -1937,8 +2065,10 @@ def test_test_iam_permissions(self): http = object() client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(RETURNED) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.test_iam_permissions(self.TABLE_REF, PERMISSIONS, timeout=7.5) - client.test_iam_permissions(self.TABLE_REF, PERMISSIONS, timeout=7.5) + final_attributes.assert_called_once_with({'path': PATH}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -1960,12 +2090,14 @@ def test_test_iam_permissions_w_invalid_table(self): with self.assertRaises(TypeError): client.test_iam_permissions(table_resource_string, PERMISSIONS) + def test_update_dataset_w_invalid_field(self): from google.cloud.bigquery.dataset import Dataset creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) with self.assertRaises(ValueError): + client.update_dataset( Dataset("{}.{}".format(self.PROJECT, self.DS_ID)), ["foo"] ) @@ -2000,11 +2132,17 @@ def test_update_dataset(self): ds.default_table_expiration_ms = EXP ds.labels = LABELS ds.access_entries = [AccessEntry("OWNER", "userByEmail", "phred@example.com")] - ds2 = client.update_dataset( - ds, - ["description", "friendly_name", "location", "labels", "access_entries"], - timeout=7.5, - ) + fields = ["description", "friendly_name", "location", "labels", "access_entries"], + + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + ds2 = client.update_dataset( + ds, + fields=fields, + timeout=7.5, + ) + final_attributes.assert_called_once_with({'path': PATH, 'fields': fields}, client, None) + + conn.api_request.assert_called_once_with( method="PATCH", data={ @@ -2046,7 +2184,11 @@ def test_update_dataset_w_custom_property(self): dataset = Dataset(DatasetReference(self.PROJECT, self.DS_ID)) dataset._properties["newAlphaProperty"] = "unreleased property" - dataset = client.update_dataset(dataset, ["newAlphaProperty"]) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + dataset = client.update_dataset(dataset, ["newAlphaProperty"]) + + final_attributes.assert_called_once_with({'path': path, 'fields':["newAlphaProperty"]}, client, None) + conn.api_request.assert_called_once_with( method="PATCH", data={"newAlphaProperty": "unreleased property"}, @@ -2093,10 +2235,14 @@ def test_update_model(self): model.friendly_name = title model.expires = expires model.labels = {"x": "y"} + fields = ["description", "friendly_name", "labels", "expires"] + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + updated_model = client.update_model( + model, fields , timeout=7.5 + ) + final_attributes.assert_called_once_with({'path': path, 'fields':fields}, client, None) + - updated_model = client.update_model( - model, ["description", "friendly_name", "labels", "expires"], timeout=7.5 - ) sent = { "description": description, @@ -2153,12 +2299,17 @@ def test_update_routine(self): routine.language = "SQL" routine.type_ = "SCALAR_FUNCTION" routine._properties["someNewField"] = "someValue" + fields = ["arguments", "language", "body", "type_", "return_type", "someNewField"], + + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + actual_routine = client.update_routine( + routine, + fields, + timeout=7.5, + ) + final_attributes.assert_called_once_with({'path': routine.path, 'fields':fields}, client, None) + - actual_routine = client.update_routine( - routine, - ["arguments", "language", "body", "type_", "return_type", "someNewField"], - timeout=7.5, - ) # TODO: routineReference isn't needed when the Routines API supports # partial updates. @@ -2177,7 +2328,11 @@ def test_update_routine(self): # ETag becomes If-Match header. routine._properties["etag"] = "im-an-etag" - client.update_routine(routine, []) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.update_routine(routine, []) + + final_attributes.assert_called_once_with({'path': routine.path, 'fields':[]}, client, None) + req = conn.api_request.call_args self.assertEqual(req[1]["headers"]["If-Match"], "im-an-etag") @@ -2228,10 +2383,14 @@ def test_update_table(self): table.description = description table.friendly_name = title table.labels = {"x": "y"} + fields = ["schema", "description", "friendly_name", "labels"] + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + updated_table = client.update_table( + table, fields, timeout=7.5 + ) + final_attributes.assert_called_once_with({'path': path, 'fields':fields}, client, None) + - updated_table = client.update_table( - table, ["schema", "description", "friendly_name", "labels"], timeout=7.5 - ) sent = { "schema": { @@ -2264,7 +2423,10 @@ def test_update_table(self): # ETag becomes If-Match header. table._properties["etag"] = "etag" - client.update_table(table, []) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.update_table(table, []) + final_attributes.assert_called_once_with({'path': path, 'fields':[]}, client, None) + req = conn.api_request.call_args self.assertEqual(req[1]["headers"]["If-Match"], "etag") @@ -2284,7 +2446,11 @@ def test_update_table_w_custom_property(self): table = Table(self.TABLE_REF) table._properties["newAlphaProperty"] = "unreleased property" - updated_table = client.update_table(table, ["newAlphaProperty"]) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + updated_table = client.update_table(table, ["newAlphaProperty"]) + + final_attributes.assert_called_once_with({'path': path, 'fields':["newAlphaProperty"]}, client, None) + conn.api_request.assert_called_once_with( method="PATCH", @@ -2312,8 +2478,11 @@ def test_update_table_only_use_legacy_sql(self): conn = client._connection = make_connection(resource) table = Table(self.TABLE_REF) table.view_use_legacy_sql = True + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + updated_table = client.update_table(table, ["view_use_legacy_sql"]) + + final_attributes.assert_called_once_with({'path': path, 'fields':["view_use_legacy_sql"]}, client, None) - updated_table = client.update_table(table, ["view_use_legacy_sql"]) conn.api_request.assert_called_once_with( method="PATCH", @@ -2376,8 +2545,11 @@ def test_update_table_w_query(self): table.view_query = query table.view_use_legacy_sql = True updated_properties = ["schema", "view_query", "expires", "view_use_legacy_sql"] + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + updated_table = client.update_table(table, updated_properties) + + final_attributes.assert_called_once_with({'path': path, 'fields':updated_properties}, client, None) - updated_table = client.update_table(table, updated_properties) self.assertEqual(updated_table.schema, table.schema) self.assertEqual(updated_table.view_query, table.view_query) @@ -2420,17 +2592,24 @@ def test_update_table_w_schema_None(self): creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(resource1, resource2) - table = client.get_table( - # Test with string for table ID - "{}.{}.{}".format( - self.TABLE_REF.project, - self.TABLE_REF.dataset_id, - self.TABLE_REF.table_id, + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + table = client.get_table( + # Test with string for table ID + "{}.{}.{}".format( + self.TABLE_REF.project, + self.TABLE_REF.dataset_id, + self.TABLE_REF.table_id, + ) ) - ) + final_attributes.assert_called_once_with({'path': path}, client, None) + + table.schema = None - updated_table = client.update_table(table, ["schema"]) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + updated_table = client.update_table(table, ["schema"]) + + final_attributes.assert_called_once_with({'path': path, 'fields':["schema"]}, client, None) self.assertEqual(len(conn.api_request.call_args_list), 2) req = conn.api_request.call_args_list[1] @@ -2460,11 +2639,20 @@ def test_update_table_delete_property(self): table = Table(self.TABLE_REF) table.description = description table.friendly_name = title - table2 = client.update_table(table, ["description", "friendly_name"]) + + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + table2 = client.update_table(table, ["description", "friendly_name"]) + + final_attributes.assert_called_once_with({'path': path, 'fields':["description", "friendly_name"]}, client, None) + self.assertEqual(table2.description, table.description) table2.description = None - table3 = client.update_table(table2, ["description"]) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + table3 = client.update_table(table2, ["description"]) + + final_attributes.assert_called_once_with({'path': path, 'fields':["description"]}, client, None) + self.assertEqual(len(conn.api_request.call_args_list), 2) req = conn.api_request.call_args_list[1] self.assertEqual(req[1]["method"], "PATCH") @@ -2480,7 +2668,11 @@ def test_list_tables_empty_w_timeout(self): conn = client._connection = make_connection({}) dataset = DatasetReference(self.PROJECT, self.DS_ID) - iterator = client.list_tables(dataset, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_tables(dataset, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path, 'page_token': None, 'max_results': None}, client, None) + self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) tables = list(page) @@ -2499,7 +2691,11 @@ def test_list_models_empty_w_timeout(self): conn = client._connection = make_connection({}) dataset_id = "{}.{}".format(self.PROJECT, self.DS_ID) - iterator = client.list_models(dataset_id, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_models(dataset_id, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path, 'page_token': iterator.next_page_token, 'max_results': None}, client, None) + page = six.next(iterator.pages) models = list(page) token = iterator.next_page_token @@ -2542,7 +2738,11 @@ def test_list_models_defaults(self): conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) - iterator = client.list_models(dataset) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_models(dataset) + + final_attributes.assert_called_once_with({'path': PATH, 'page_token':TOKEN , 'max_results': None}, client, None) + self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) models = list(page) @@ -2568,8 +2768,13 @@ def test_list_routines_empty_w_timeout(self): creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection({}) + path = "/projects/test-routines/datasets/test_routines/routines" + + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_routines("test-routines.test_routines", timeout=7.5) + + final_attributes.assert_called_once_with({'path': path, 'page_token' : None, 'max_results': None}, client, None) - iterator = client.list_routines("test-routines.test_routines", timeout=7.5) page = six.next(iterator.pages) routines = list(page) token = iterator.next_page_token @@ -2578,7 +2783,7 @@ def test_list_routines_empty_w_timeout(self): self.assertIsNone(token) conn.api_request.assert_called_once_with( method="GET", - path="/projects/test-routines/datasets/test_routines/routines", + path=path, query_params={}, timeout=7.5, ) @@ -2616,8 +2821,11 @@ def test_list_routines_defaults(self): client = self._make_one(project=project_id, credentials=creds) conn = client._connection = make_connection(resource) dataset = DatasetReference(client.project, dataset_id) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_routines(dataset) + + final_attributes.assert_called_once_with({'path': path, 'page_token': None, 'max_results': None}, client, None) - iterator = client.list_routines(dataset) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) routines = list(page) @@ -2680,8 +2888,11 @@ def test_list_tables_defaults(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_tables(dataset) + + final_attributes.assert_called_once_with({'path': PATH, 'page_token': None, 'max_results': None}, client, None) - iterator = client.list_tables(dataset) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) tables = list(page) @@ -2734,13 +2945,15 @@ def test_list_tables_explicit(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + iterator = client.list_tables( + # Test with string for dataset ID. + self.DS_ID, + max_results=3, + page_token=TOKEN, + ) + final_attributes.assert_called_once_with({'path': PATH, 'page_token': TOKEN, 'max_results': 3}, client, None) - iterator = client.list_tables( - # Test with string for dataset ID. - self.DS_ID, - max_results=3, - page_token=TOKEN, - ) self.assertEqual(iterator.dataset, dataset) page = six.next(iterator.pages) tables = list(page) @@ -2777,7 +2990,11 @@ def test_delete_dataset(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(*([{}] * len(datasets))) for arg in datasets: - client.delete_dataset(arg, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_dataset(arg, timeout=7.5) + + final_attributes.assert_called_once_with({'path': PATH}, client, None) + conn.api_request.assert_called_with( method="DELETE", path="/%s" % PATH, query_params={}, timeout=7.5 ) @@ -2791,7 +3008,10 @@ def test_delete_dataset_delete_contents(self): conn = client._connection = make_connection({}, {}) ds_ref = DatasetReference(self.PROJECT, self.DS_ID) for arg in (ds_ref, Dataset(ds_ref)): - client.delete_dataset(arg, delete_contents=True) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_dataset(arg, delete_contents=True) + + final_attributes.assert_called_once_with({'path': PATH}, client, None) conn.api_request.assert_called_with( method="DELETE", path="/%s" % PATH, @@ -2817,7 +3037,10 @@ def test_delete_dataset_w_not_found_ok_false(self): ) with self.assertRaises(google.api_core.exceptions.NotFound): - client.delete_dataset(self.DS_ID) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_dataset(self.DS_ID) + + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_with( method="DELETE", path=path, query_params={}, timeout=None @@ -2832,7 +3055,10 @@ def test_delete_dataset_w_not_found_ok_true(self): google.api_core.exceptions.NotFound("dataset not found") ) - client.delete_dataset(self.DS_ID, not_found_ok=True) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_dataset(self.DS_ID, not_found_ok=True) + + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_with( method="DELETE", path=path, query_params={}, timeout=None @@ -2858,7 +3084,10 @@ def test_delete_model(self): conn = client._connection = make_connection(*([{}] * len(models))) for arg in models: - client.delete_model(arg, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_model(arg, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path, 'not_found_okay': False}, client, None) conn.api_request.assert_called_with( method="DELETE", path="/%s" % path, timeout=7.5 ) @@ -2895,10 +3124,11 @@ def test_delete_model_w_not_found_ok_true(self): conn = client._connection = make_connection( google.api_core.exceptions.NotFound("model not found") ) - - client.delete_model( - "{}.{}".format(self.DS_ID, self.MODEL_ID), not_found_ok=True - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_model( + "{}.{}".format(self.DS_ID, self.MODEL_ID), not_found_ok=True + ) + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) @@ -2914,14 +3144,19 @@ def test_delete_routine(self): ] creds = _make_credentials() http = object() + path = "/projects/test-routine-project/datasets/test_routines/routines/minimal_routine" client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(*([{}] * len(routines))) for routine in routines: - client.delete_routine(routine, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_routine(routine, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path}, client, None) + conn.api_request.assert_called_with( method="DELETE", - path="/projects/test-routine-project/datasets/test_routines/routines/minimal_routine", + path=path, timeout=7.5, ) @@ -2938,13 +3173,17 @@ def test_delete_routine_w_not_found_ok_false(self): conn = client._connection = make_connection( google.api_core.exceptions.NotFound("routine not found") ) + path = "/projects/routines-project/datasets/test_routines/routines/test_routine", with self.assertRaises(google.api_core.exceptions.NotFound): - client.delete_routine("routines-project.test_routines.test_routine") + with mock.patch( 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_routine("routines-project.test_routines.test_routine") + + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_with( method="DELETE", - path="/projects/routines-project/datasets/test_routines/routines/test_routine", + path=path, timeout=None, ) @@ -2955,14 +3194,17 @@ def test_delete_routine_w_not_found_ok_true(self): conn = client._connection = make_connection( google.api_core.exceptions.NotFound("routine not found") ) + path = "/projects/routines-project/datasets/test_routines/routines/test_routine", - client.delete_routine( - "routines-project.test_routines.test_routine", not_found_ok=True - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_routine( + "routines-project.test_routines.test_routine", not_found_ok=True + ) + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_with( method="DELETE", - path="/projects/routines-project/datasets/test_routines/routines/test_routine", + path=path, timeout=None, ) @@ -2989,7 +3231,11 @@ def test_delete_table(self): conn = client._connection = make_connection(*([{}] * len(tables))) for arg in tables: - client.delete_table(arg, timeout=7.5) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_table(arg, timeout=7.5) + + final_attributes.assert_called_once_with({'path': path}, client, None) + conn.api_request.assert_called_with( method="DELETE", path="/%s" % path, timeout=7.5 ) @@ -3012,7 +3258,10 @@ def test_delete_table_w_not_found_ok_false(self): ) with self.assertRaises(google.api_core.exceptions.NotFound): - client.delete_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) + + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) @@ -3027,9 +3276,12 @@ def test_delete_table_w_not_found_ok_true(self): google.api_core.exceptions.NotFound("table not found") ) - client.delete_table( - "{}.{}".format(self.DS_ID, self.TABLE_ID), not_found_ok=True - ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + client.delete_table( + "{}.{}".format(self.DS_ID, self.TABLE_ID), not_found_ok=True + ) + + final_attributes.assert_called_once_with({'path': path}, client, None) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) From ac56488c01b7b2ea10dd32c256557f328fad4fb7 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Sat, 15 Aug 2020 11:14:27 -0400 Subject: [PATCH 19/41] adding reamining client tests and all job tests --- .../cloud/bigquery/opentelemetry_tracing.py | 5 - tests/unit/test_client.py | 285 +++++++++--------- tests/unit/test_job.py | 282 +++++++++++------ 3 files changed, 340 insertions(+), 232 deletions(-) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 8fa8ceb65..530cbed3c 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -42,7 +42,6 @@ tracer = trace.get_tracer(__name__) - @contextmanager def create_span(name, attributes=None, client=None, job_ref=None): # yield new span value @@ -88,7 +87,6 @@ def _get_final_span_attributes(attributes=None, client=None, job_ref=None): OpenTelemetry. """ - if client: client_attributes = _set_client_attributes(client) _default_attributes.update(client_attributes) @@ -120,6 +118,3 @@ def _set_job_attributes(job_ref): "errorResult": job_ref.error_result, "state": job_ref.state, } - - - diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index a4ba43e76..88b49fa08 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -90,7 +90,6 @@ def _make_list_partitons_meta_info(project, dataset_id, table_id, num_rows=0): class TestClient(unittest.TestCase): - PROJECT = "PROJECT" DS_ID = "DATASET_ID" TABLE_ID = "TABLE_ID" @@ -249,7 +248,8 @@ def test__get_query_results_miss_w_explicit_project_and_timeout(self): conn = client._connection = make_connection() path = "/projects/other-project/queries/nothere" with self.assertRaises(NotFound): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client._get_query_results( "nothere", None, @@ -260,7 +260,6 @@ def test__get_query_results_miss_w_explicit_project_and_timeout(self): ) final_attributes.assert_called_once_with({'path': path}, client, None) - conn.api_request.assert_called_once_with( method="GET", path=path, @@ -361,7 +360,8 @@ def test_get_service_account_email_w_custom_retry(self): ) with api_request_patcher as fake_api_request: - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: service_account_email = client.get_service_account_email( retry=retry, timeout=7.5 ) @@ -504,7 +504,7 @@ def test_list_datasets_defaults(self): conn = client._connection = make_connection(DATA) with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: iterator = client.list_datasets() - final_attributes.assert_called_once_with({'path': PATH, 'page_token':TOKEN}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': None}, client, None) page = six.next(iterator.pages) datasets = list(page) @@ -553,7 +553,6 @@ def test_list_datasets_explicit_response_missing_datasets_key(self): ) final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': TOKEN}, client, None) - page = six.next(iterator.pages) datasets = list(page) token = iterator.next_page_token @@ -588,8 +587,8 @@ def test_dataset_with_specified_project(self): warning for warning in warned if warning.category in (DeprecationWarning, PendingDeprecationWarning) - and "Client.dataset" in str(warning) - and "my_project.my_dataset" in str(warning) + and "Client.dataset" in str(warning) + and "my_project.my_dataset" in str(warning) ] assert matches, "A Client.dataset deprecation warning was not raised." self.assertIsInstance(dataset, DatasetReference) @@ -611,8 +610,8 @@ def test_dataset_with_default_project(self): warning for warning in warned if warning.category in (DeprecationWarning, PendingDeprecationWarning) - and "Client.dataset" in str(warning) - and "my_project.my_dataset" in str(warning) + and "Client.dataset" in str(warning) + and "my_project.my_dataset" in str(warning) ] assert matches, "A Client.dataset deprecation warning was not raised." self.assertIsInstance(dataset, DatasetReference) @@ -637,7 +636,6 @@ def test_get_dataset(self): final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) - conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 ) @@ -648,25 +646,26 @@ def test_get_dataset(self): # Not a cloud API exception (missing 'errors' field). client._connection = make_connection(Exception(""), resource) with self.assertRaises(Exception): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with({'path': path}, client, None) - + final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) # Zero-length errors field. client._connection = make_connection(ServerError(""), resource) with self.assertRaises(ServerError): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with({'path': path}, client, None) - + final_attributes.assert_called_once_with({'path': '/%s' % path}, client, None) # Non-retryable reason. client._connection = make_connection( ServerError("", errors=[{"reason": "serious"}]), resource ) with self.assertRaises(ServerError): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.get_dataset(dataset_ref) final_attributes.assert_called_once_with({'path': path}, client, None) @@ -676,10 +675,11 @@ def test_get_dataset(self): ServerError("", errors=[{"reason": "backendError"}]), resource ) with self.assertRaises(ServerError): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.get_dataset(dataset_ref, retry=None) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({'path': '/%s' % path}, client, None) # Retryable reason, default retry: success. client._connection = make_connection( @@ -690,7 +690,7 @@ def test_get_dataset(self): # Test with a string for dataset ID. dataset_ref.dataset_id ) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({'path': '/%s' % path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) @@ -705,7 +705,7 @@ def test_create_bqstorage_client(self): client = self._make_one(project=self.PROJECT, credentials=creds) with mock.patch( - "google.cloud.bigquery_storage_v1.BigQueryReadClient", mock_client + "google.cloud.bigquery_storage_v1.BigQueryReadClient", mock_client ): bqstorage_client = client._create_bqstorage_client() @@ -719,7 +719,7 @@ def test_create_bqstorage_client_missing_dependency(self): def fail_bqstorage_import(name, globals, locals, fromlist, level): # NOTE: *very* simplified, assuming a straightforward absolute import return "bigquery_storage_v1" in name or ( - fromlist is not None and "bigquery_storage_v1" in fromlist + fromlist is not None and "bigquery_storage_v1" in fromlist ) no_bqstorage = maybe_fail_import(predicate=fail_bqstorage_import) @@ -732,7 +732,7 @@ def fail_bqstorage_import(name, globals, locals, fromlist, level): warning for warning in warned if "not installed" in str(warning) - and "google-cloud-bigquery-storage" in str(warning) + and "google-cloud-bigquery-storage" in str(warning) ] assert matching_warnings, "Missing dependency warning not raised." @@ -756,7 +756,6 @@ def test_create_dataset_minimal(self): final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) - self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) self.assertEqual(after.etag, RESOURCE["etag"]) @@ -912,7 +911,7 @@ def test_create_dataset_w_client_location_wo_dataset_location(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: after = client.create_dataset(before) - final_attributes.assert_called_once_with({'path': PATH}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -956,7 +955,7 @@ def test_create_dataset_w_client_location_w_dataset_location(self): before.location = OTHER_LOCATION with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: after = client.create_dataset(before) - final_attributes.assert_called_once_with({'path': PATH}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -1102,12 +1101,12 @@ def test_create_dataset_alreadyexists_w_exists_ok_false(self): ) with pytest.raises(google.api_core.exceptions.AlreadyExists): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.create_dataset(self.DS_ID) final_attributes.assert_called_once_with({'path': 'path'}, client, None) - def test_create_dataset_alreadyexists_w_exists_ok_true(self): post_path = "/projects/{}/datasets".format(self.PROJECT) get_path = "/projects/{}/datasets/{}".format(self.PROJECT, self.DS_ID) @@ -1126,7 +1125,7 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(self): ) with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: dataset = client.create_dataset(self.DS_ID, exists_ok=True) - final_attributes.assert_called_once_with({'path': post_path}, client, None) + final_attributes.assert_called_with({'path': get_path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -1173,12 +1172,11 @@ def test_create_routine_w_minimal_resource(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: actual_routine = client.create_routine(routine, timeout=7.5) - final_attributes.assert_called_once_with({'path': path, 'exists_okay': False}, client, None) - + final_attributes.assert_called_once_with({'path': path, 'exists_ok': False}, client, None) conn.api_request.assert_called_once_with( method="POST", - path='path', + path=path, data=resource, timeout=7.5, ) @@ -1199,7 +1197,8 @@ def test_create_routine_w_conflict(self): routine = Routine(full_routine_id) with pytest.raises(google.api_core.exceptions.AlreadyExists): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.create_routine(routine) final_attributes.assert_called_once_with({'path': path, 'exists_okay': False}, client, None) @@ -1230,7 +1229,7 @@ def test_create_routine_w_conflict_exists_ok(self): "routineId": "minimal_routine", } } - path = "/projects/test-routine-project/datasets/test_routines/routines", + path = "/projects/test-routine-project/datasets/test_routines/routines" conn = client._connection = make_connection( google.api_core.exceptions.AlreadyExists("routine already exists"), resource @@ -1240,8 +1239,7 @@ def test_create_routine_w_conflict_exists_ok(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: actual_routine = client.create_routine(routine, exists_ok=True) - final_attributes.assert_called_once_with({'path': path, 'exists_okay': True}, client, None) - + final_attributes.assert_called_with({'path': path + '/minimal_routine'}, client, None) self.assertEqual(actual_routine.project, "test-routine-project") self.assertEqual(actual_routine.dataset_id, "test_routines") @@ -1276,7 +1274,7 @@ def test_create_table_w_day_partition(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.create_table(table, timeout=7.5) - final_attributes.assert_called_once_with({'path': path, 'dataset_id' : table.dataset_id, 'exists_okay': True}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1310,7 +1308,7 @@ def test_create_table_w_custom_property(self): table._properties["newAlphaProperty"] = "unreleased property" with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1347,7 +1345,7 @@ def test_create_table_w_encryption_configuration(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path' : path, 'dataset_id' : table.dataset_id}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1379,7 +1377,7 @@ def test_create_table_w_day_partition_and_expire(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1440,8 +1438,7 @@ def test_create_table_w_schema_and_query(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': path, 'dataset_id': table.dataset_id}, client, None) - + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1505,7 +1502,7 @@ def test_create_table_w_external(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1542,7 +1539,8 @@ def test_create_table_w_reference(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.create_table(self.TABLE_REF) - final_attributes.assert_called_once_with({'path': path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': self.TABLE_REF.dataset_id}, + client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1569,7 +1567,8 @@ def test_create_table_w_fully_qualified_string(self): got = client.create_table( "{}.{}.{}".format(self.PROJECT, self.DS_ID, self.TABLE_ID) ) - final_attributes.assert_called_once_with({'path': path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': self.TABLE_REF.dataset_id}, + client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1595,7 +1594,8 @@ def test_create_table_w_string(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_once_with({'path': path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': self.TABLE_REF.dataset_id}, + client, None) conn.api_request.assert_called_once_with( method="POST", @@ -1623,11 +1623,12 @@ def test_create_table_alreadyexists_w_exists_ok_false(self): ) with pytest.raises(google.api_core.exceptions.AlreadyExists): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_with({'path': post_path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) - + final_attributes.assert_called_with({'path': post_path, 'dataset_id': self.TABLE_REF.dataset_id}, client, + None) conn.api_request.assert_called_once_with( method="POST", @@ -1661,9 +1662,7 @@ def test_create_table_alreadyexists_w_exists_ok_true(self): got = client.create_table( "{}.{}".format(self.DS_ID, self.TABLE_ID), exists_ok=True ) - final_attributes.assert_called_once_with({'path': post_path, 'dataset_id': self.TABLE_REF.dataset_id}, client, None) - - + final_attributes.assert_called_with({'path': get_path}, client, None) self.assertEqual(got.project, self.PROJECT) self.assertEqual(got.dataset_id, self.DS_ID) @@ -1721,8 +1720,7 @@ def test_get_model(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.get_model(model_ref, timeout=7.5) - final_attributes.assert_called_once_with({'path': path}, client, None) - + final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 @@ -1751,7 +1749,7 @@ def test_get_model_w_string(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: got = client.get_model(model_id) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=None @@ -1779,17 +1777,17 @@ def test_get_routine(self): }, "routineType": "SCALAR_FUNCTION", } - path = "/projects/test-routine-project/datasets/test_routines/routines/minimal_routine", + path = "/projects/test-routine-project/datasets/test_routines/routines/minimal_routine" client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: actual_routine = client.get_routine(routine, timeout=7.5) final_attributes.assert_called_once_with({'path': path}, client, None) - conn.api_request.assert_called_once_with( method="GET", path=path, @@ -1825,7 +1823,7 @@ def test_get_table(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: table = client.get_table(self.TABLE_REF, timeout=7.5) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 @@ -1848,6 +1846,7 @@ def test_get_table_sets_user_agent(self): client_info=user_agent_override, _http=http, ) + client.get_table(self.TABLE_REF) expected_user_agent = user_agent_override.to_user_agent() http.request.assert_called_once_with( @@ -1902,7 +1901,7 @@ def test_get_iam_policy(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: policy = client.get_iam_policy(self.TABLE_REF, timeout=7.5) - final_attributes.assert_called_once_with({'path': PATH}, client, None) + final_attributes.assert_called_once_with({'path': PATH, 'body': BODY}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -1925,7 +1924,6 @@ def test_get_iam_policy_w_invalid_table(self): with self.assertRaises(TypeError): client.get_iam_policy(table_resource_string) - def test_get_iam_policy_w_invalid_version(self): creds = _make_credentials() http = object() @@ -2022,11 +2020,12 @@ def test_set_iam_policy_invalid_policy(self): client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) with self.assertRaises(TypeError): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.set_iam_policy(self.TABLE_REF, invalid_policy_repr) - final_attributes.assert_called_once_with({'path': "{}:setIamPolicy".format(self.TABLE_REF.path)}, client, None) - + final_attributes.assert_called_once_with({'path': "{}:setIamPolicy".format(self.TABLE_REF.path)}, client, + None) def test_set_iam_policy_w_invalid_table(self): from google.api_core.iam import Policy @@ -2044,12 +2043,12 @@ def test_set_iam_policy_w_invalid_table(self): ) with self.assertRaises(TypeError): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.set_iam_policy(table_resource_string, policy) final_attributes.assert_called_once_with({'path': table_resource_string}, client, None) - def test_test_iam_permissions(self): PATH = "/projects/%s/datasets/%s/tables/%s:testIamPermissions" % ( self.PROJECT, @@ -2068,7 +2067,7 @@ def test_test_iam_permissions(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.test_iam_permissions(self.TABLE_REF, PERMISSIONS, timeout=7.5) - final_attributes.assert_called_once_with({'path': PATH}, client, None) + final_attributes.assert_called_once_with({'path': PATH, 'body': BODY}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -2090,14 +2089,12 @@ def test_test_iam_permissions_w_invalid_table(self): with self.assertRaises(TypeError): client.test_iam_permissions(table_resource_string, PERMISSIONS) - def test_update_dataset_w_invalid_field(self): from google.cloud.bigquery.dataset import Dataset creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) with self.assertRaises(ValueError): - client.update_dataset( Dataset("{}.{}".format(self.PROJECT, self.DS_ID)), ["foo"] ) @@ -2132,7 +2129,7 @@ def test_update_dataset(self): ds.default_table_expiration_ms = EXP ds.labels = LABELS ds.access_entries = [AccessEntry("OWNER", "userByEmail", "phred@example.com")] - fields = ["description", "friendly_name", "location", "labels", "access_entries"], + fields = ["description", "friendly_name", "location", "labels", "access_entries"] with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: ds2 = client.update_dataset( @@ -2140,8 +2137,7 @@ def test_update_dataset(self): fields=fields, timeout=7.5, ) - final_attributes.assert_called_once_with({'path': PATH, 'fields': fields}, client, None) - + final_attributes.assert_called_once_with({'path': "/" + PATH, 'fields': fields}, client, None) conn.api_request.assert_called_once_with( method="PATCH", @@ -2187,7 +2183,7 @@ def test_update_dataset_w_custom_property(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: dataset = client.update_dataset(dataset, ["newAlphaProperty"]) - final_attributes.assert_called_once_with({'path': path, 'fields':["newAlphaProperty"]}, client, None) + final_attributes.assert_called_once_with({'path': path, 'fields': ["newAlphaProperty"]}, client, None) conn.api_request.assert_called_once_with( method="PATCH", @@ -2238,11 +2234,9 @@ def test_update_model(self): fields = ["description", "friendly_name", "labels", "expires"] with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: updated_model = client.update_model( - model, fields , timeout=7.5 + model, fields, timeout=7.5 ) - final_attributes.assert_called_once_with({'path': path, 'fields':fields}, client, None) - - + final_attributes.assert_called_once_with({'path': "/" + path, 'fields': fields}, client, None) sent = { "description": description, @@ -2299,7 +2293,7 @@ def test_update_routine(self): routine.language = "SQL" routine.type_ = "SCALAR_FUNCTION" routine._properties["someNewField"] = "someValue" - fields = ["arguments", "language", "body", "type_", "return_type", "someNewField"], + fields = ["arguments", "language", "body", "type_", "return_type", "someNewField"] with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: actual_routine = client.update_routine( @@ -2307,9 +2301,7 @@ def test_update_routine(self): fields, timeout=7.5, ) - final_attributes.assert_called_once_with({'path': routine.path, 'fields':fields}, client, None) - - + final_attributes.assert_called_once_with({'path': routine.path, 'fields': fields}, client, None) # TODO: routineReference isn't needed when the Routines API supports # partial updates. @@ -2331,7 +2323,7 @@ def test_update_routine(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.update_routine(routine, []) - final_attributes.assert_called_once_with({'path': routine.path, 'fields':[]}, client, None) + final_attributes.assert_called_once_with({'path': routine.path, 'fields': []}, client, None) req = conn.api_request.call_args self.assertEqual(req[1]["headers"]["If-Match"], "im-an-etag") @@ -2388,9 +2380,8 @@ def test_update_table(self): updated_table = client.update_table( table, fields, timeout=7.5 ) - final_attributes.assert_called_once_with({'path': path, 'fields':fields}, client, None) - - + span_path = '/%s' % path + final_attributes.assert_called_once_with({'path': span_path, 'fields': fields}, client, None) sent = { "schema": { @@ -2425,7 +2416,7 @@ def test_update_table(self): table._properties["etag"] = "etag" with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.update_table(table, []) - final_attributes.assert_called_once_with({'path': path, 'fields':[]}, client, None) + final_attributes.assert_called_once_with({'path': "/" + path, 'fields': []}, client, None) req = conn.api_request.call_args self.assertEqual(req[1]["headers"]["If-Match"], "etag") @@ -2449,8 +2440,7 @@ def test_update_table_w_custom_property(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: updated_table = client.update_table(table, ["newAlphaProperty"]) - final_attributes.assert_called_once_with({'path': path, 'fields':["newAlphaProperty"]}, client, None) - + final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["newAlphaProperty"]}, client, None) conn.api_request.assert_called_once_with( method="PATCH", @@ -2481,8 +2471,8 @@ def test_update_table_only_use_legacy_sql(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: updated_table = client.update_table(table, ["view_use_legacy_sql"]) - final_attributes.assert_called_once_with({'path': path, 'fields':["view_use_legacy_sql"]}, client, None) - + final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["view_use_legacy_sql"]}, client, + None) conn.api_request.assert_called_once_with( method="PATCH", @@ -2548,8 +2538,7 @@ def test_update_table_w_query(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: updated_table = client.update_table(table, updated_properties) - final_attributes.assert_called_once_with({'path': path, 'fields':updated_properties}, client, None) - + final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': updated_properties}, client, None) self.assertEqual(updated_table.schema, table.schema) self.assertEqual(updated_table.view_query, table.view_query) @@ -2601,15 +2590,14 @@ def test_update_table_w_schema_None(self): self.TABLE_REF.table_id, ) ) - final_attributes.assert_called_once_with({'path': path}, client, None) - + final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) table.schema = None with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: updated_table = client.update_table(table, ["schema"]) - final_attributes.assert_called_once_with({'path': path, 'fields':["schema"]}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["schema"]}, client, None) self.assertEqual(len(conn.api_request.call_args_list), 2) req = conn.api_request.call_args_list[1] @@ -2643,7 +2631,8 @@ def test_update_table_delete_property(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: table2 = client.update_table(table, ["description", "friendly_name"]) - final_attributes.assert_called_once_with({'path': path, 'fields':["description", "friendly_name"]}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["description", "friendly_name"]}, + client, None) self.assertEqual(table2.description, table.description) table2.description = None @@ -2651,7 +2640,7 @@ def test_update_table_delete_property(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: table3 = client.update_table(table2, ["description"]) - final_attributes.assert_called_once_with({'path': path, 'fields':["description"]}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["description"]}, client, None) self.assertEqual(len(conn.api_request.call_args_list), 2) req = conn.api_request.call_args_list[1] @@ -2694,7 +2683,8 @@ def test_list_models_empty_w_timeout(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: iterator = client.list_models(dataset_id, timeout=7.5) - final_attributes.assert_called_once_with({'path': path, 'page_token': iterator.next_page_token, 'max_results': None}, client, None) + final_attributes.assert_called_once_with( + {'path': path, 'page_token': iterator.next_page_token, 'max_results': None}, client, None) page = six.next(iterator.pages) models = list(page) @@ -2741,7 +2731,8 @@ def test_list_models_defaults(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: iterator = client.list_models(dataset) - final_attributes.assert_called_once_with({'path': PATH, 'page_token':TOKEN , 'max_results': None}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': None, 'max_results': None}, + client, None) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2773,7 +2764,7 @@ def test_list_routines_empty_w_timeout(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: iterator = client.list_routines("test-routines.test_routines", timeout=7.5) - final_attributes.assert_called_once_with({'path': path, 'page_token' : None, 'max_results': None}, client, None) + final_attributes.assert_called_once_with({'path': path, 'page_token': None, 'max_results': None}, client, None) page = six.next(iterator.pages) routines = list(page) @@ -2891,7 +2882,8 @@ def test_list_tables_defaults(self): with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: iterator = client.list_tables(dataset) - final_attributes.assert_called_once_with({'path': PATH, 'page_token': None, 'max_results': None}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': None, 'max_results': None}, + client, None) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2952,7 +2944,8 @@ def test_list_tables_explicit(self): max_results=3, page_token=TOKEN, ) - final_attributes.assert_called_once_with({'path': PATH, 'page_token': TOKEN, 'max_results': 3}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': TOKEN, 'max_results': 3}, client, + None) self.assertEqual(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2990,10 +2983,11 @@ def test_delete_dataset(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(*([{}] * len(datasets))) for arg in datasets: - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_dataset(arg, timeout=7.5) - final_attributes.assert_called_once_with({'path': PATH}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) conn.api_request.assert_called_with( method="DELETE", path="/%s" % PATH, query_params={}, timeout=7.5 @@ -3008,10 +3002,11 @@ def test_delete_dataset_delete_contents(self): conn = client._connection = make_connection({}, {}) ds_ref = DatasetReference(self.PROJECT, self.DS_ID) for arg in (ds_ref, Dataset(ds_ref)): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_dataset(arg, delete_contents=True) - final_attributes.assert_called_once_with({'path': PATH}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) conn.api_request.assert_called_with( method="DELETE", path="/%s" % PATH, @@ -3037,7 +3032,8 @@ def test_delete_dataset_w_not_found_ok_false(self): ) with self.assertRaises(google.api_core.exceptions.NotFound): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_dataset(self.DS_ID) final_attributes.assert_called_once_with({'path': path}, client, None) @@ -3084,10 +3080,11 @@ def test_delete_model(self): conn = client._connection = make_connection(*([{}] * len(models))) for arg in models: - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_model(arg, timeout=7.5) - final_attributes.assert_called_once_with({'path': path, 'not_found_okay': False}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path, 'not_found_okay': False}, client, None) conn.api_request.assert_called_with( method="DELETE", path="/%s" % path, timeout=7.5 ) @@ -3128,7 +3125,7 @@ def test_delete_model_w_not_found_ok_true(self): client.delete_model( "{}.{}".format(self.DS_ID, self.MODEL_ID), not_found_ok=True ) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({'path': path, 'not_found_okay': True}, client, None) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) @@ -3149,7 +3146,8 @@ def test_delete_routine(self): conn = client._connection = make_connection(*([{}] * len(routines))) for routine in routines: - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_routine(routine, timeout=7.5) final_attributes.assert_called_once_with({'path': path}, client, None) @@ -3173,10 +3171,11 @@ def test_delete_routine_w_not_found_ok_false(self): conn = client._connection = make_connection( google.api_core.exceptions.NotFound("routine not found") ) - path = "/projects/routines-project/datasets/test_routines/routines/test_routine", + path = "/projects/routines-project/datasets/test_routines/routines/test_routine" with self.assertRaises(google.api_core.exceptions.NotFound): - with mock.patch( 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_routine("routines-project.test_routines.test_routine") final_attributes.assert_called_once_with({'path': path}, client, None) @@ -3194,7 +3193,7 @@ def test_delete_routine_w_not_found_ok_true(self): conn = client._connection = make_connection( google.api_core.exceptions.NotFound("routine not found") ) - path = "/projects/routines-project/datasets/test_routines/routines/test_routine", + path = "/projects/routines-project/datasets/test_routines/routines/test_routine" with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_routine( @@ -3231,10 +3230,11 @@ def test_delete_table(self): conn = client._connection = make_connection(*([{}] * len(tables))) for arg in tables: - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_table(arg, timeout=7.5) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) conn.api_request.assert_called_with( method="DELETE", path="/%s" % path, timeout=7.5 @@ -3258,7 +3258,8 @@ def test_delete_table_w_not_found_ok_false(self): ) with self.assertRaises(google.api_core.exceptions.NotFound): - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: client.delete_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) final_attributes.assert_called_once_with({'path': path}, client, None) @@ -4128,9 +4129,9 @@ def _initiate_resumable_upload_helper(self, num_retries=None): # Check the returned values. self.assertIsInstance(upload, ResumableUpload) upload_url = ( - "https://bigquery.googleapis.com/upload/bigquery/v2/projects/" - + self.PROJECT - + "/jobs?uploadType=resumable" + "https://bigquery.googleapis.com/upload/bigquery/v2/projects/" + + self.PROJECT + + "/jobs?uploadType=resumable" ) self.assertEqual(upload.upload_url, upload_url) expected_headers = _get_upload_headers(conn.user_agent) @@ -4197,20 +4198,20 @@ def _do_multipart_upload_success_helper(self, get_boundary, num_retries=None): get_boundary.assert_called_once_with() upload_url = ( - "https://bigquery.googleapis.com/upload/bigquery/v2/projects/" - + self.PROJECT - + "/jobs?uploadType=multipart" + "https://bigquery.googleapis.com/upload/bigquery/v2/projects/" + + self.PROJECT + + "/jobs?uploadType=multipart" ) payload = ( - b"--==0==\r\n" - + b"content-type: application/json; charset=UTF-8\r\n\r\n" - + json.dumps(metadata).encode("utf-8") - + b"\r\n" - + b"--==0==\r\n" - + b"content-type: */*\r\n\r\n" - + data - + b"\r\n" - + b"--==0==--" + b"--==0==\r\n" + + b"content-type: application/json; charset=UTF-8\r\n\r\n" + + json.dumps(metadata).encode("utf-8") + + b"\r\n" + + b"--==0==\r\n" + + b"content-type: */*\r\n\r\n" + + data + + b"\r\n" + + b"--==0==--" ) headers = _get_upload_headers(conn.user_agent) headers["content-type"] = b'multipart/related; boundary="==0=="' @@ -5022,7 +5023,7 @@ def test_query_preserving_explicit_job_config(self): from google.cloud.bigquery import QueryJobConfig - client = self._make_one(project=self.PROJECT, credentials=creds, _http=http,) + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http, ) conn = client._connection = make_connection(resource) job_config = QueryJobConfig() @@ -6039,7 +6040,7 @@ def test_insert_rows_from_dataframe(self): actual_calls = conn.api_request.call_args_list for call, expected_data in six.moves.zip_longest( - actual_calls, EXPECTED_SENT_DATA + actual_calls, EXPECTED_SENT_DATA ): expected_call = mock.call( method="POST", path=API_PATH, data=expected_data, timeout=7.5 @@ -6107,7 +6108,7 @@ def test_insert_rows_from_dataframe_nan(self): actual_calls = conn.api_request.call_args_list for call, expected_data in six.moves.zip_longest( - actual_calls, EXPECTED_SENT_DATA + actual_calls, EXPECTED_SENT_DATA ): expected_call = mock.call( method="POST", path=API_PATH, data=expected_data, timeout=7.5 @@ -7068,7 +7069,7 @@ def test_load_table_from_file_with_rewind(self): file_obj.seek(2) with self._make_do_upload_patch( - client, "_do_resumable_upload", self.EXPECTED_CONFIGURATION + client, "_do_resumable_upload", self.EXPECTED_CONFIGURATION ): client.load_table_from_file(file_obj, self.TABLE_REF, rewind=True) @@ -7526,8 +7527,8 @@ def test_load_table_from_dataframe_no_schema_warning_wo_pyarrow(self): warning for warning in warned if warning.category in (DeprecationWarning, PendingDeprecationWarning) - and "could not be detected" in str(warning) - and "please provide a schema" in str(warning) + and "could not be detected" in str(warning) + and "please provide a schema" in str(warning) ] assert matches, "A missing schema deprecation warning was not raised." @@ -7803,7 +7804,7 @@ def test_load_table_from_dataframe_w_partial_schema_extra_types(self): ) job_config = job.LoadJobConfig(schema=schema) with load_patch as load_table_from_file, pytest.raises( - ValueError + ValueError ) as exc_context: client.load_table_from_dataframe( dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION @@ -7839,7 +7840,7 @@ def test_load_table_from_dataframe_w_partial_schema_missing_types(self): schema = (SchemaField("string_col", "STRING"),) job_config = job.LoadJobConfig(schema=schema) with pyarrow_patch, load_patch as load_table_from_file, warnings.catch_warnings( - record=True + record=True ) as warned: client.load_table_from_dataframe( dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION @@ -7888,7 +7889,7 @@ def test_load_table_from_dataframe_w_schema_wo_pyarrow(self): pyarrow_patch = mock.patch("google.cloud.bigquery.client.pyarrow", None) with load_patch as load_table_from_file, pyarrow_patch, warnings.catch_warnings( - record=True + record=True ) as warned: client.load_table_from_dataframe( dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 9cd3631e1..4a8bdd467 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -71,17 +71,17 @@ def _make_connection(*responses): def _make_job_resource( - creation_time_ms=1437767599006, - started_time_ms=1437767600007, - ended_time_ms=1437767601008, - started=False, - ended=False, - etag="abc-def-hjk", - endpoint="https://bigquery.googleapis.com", - job_type="load", - job_id="a-random-id", - project_id="some-project", - user_email="bq-user@example.com", + creation_time_ms=1437767599006, + started_time_ms=1437767600007, + ended_time_ms=1437767601008, + started=False, + ended=False, + etag="abc-def-hjk", + endpoint="https://bigquery.googleapis.com", + job_type="load", + job_id="a-random-id", + project_id="some-project", + user_email="bq-user@example.com", ): resource = { "configuration": {job_type: {}}, @@ -620,13 +620,16 @@ def test__begin_defaults(self): builder.return_value = resource 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() - job._begin() + final_attributes.assert_called_with({'path': path}, None, job) call_api.assert_called_once_with( DEFAULT_RETRY, method="POST", - path="/projects/{}/jobs".format(self.PROJECT), + path=path, data=resource, timeout=None, ) @@ -651,13 +654,16 @@ def test__begin_explicit(self): call_api = client._call_api = mock.Mock() 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) - job._begin(client=client, retry=retry, timeout=7.5) + final_attributes.assert_called_with({'path': path}, None, job) call_api.assert_called_once_with( retry, method="POST", - path="/projects/{}/jobs".format(self.PROJECT), + path=path, data=resource, timeout=7.5, ) @@ -672,7 +678,11 @@ def test_exists_defaults_miss(self): call_api = job._client._call_api = mock.Mock() call_api.side_effect = NotFound("testing") - self.assertFalse(job.exists()) + 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) call_api.assert_called_once_with( DEFAULT_RETRY, @@ -699,8 +709,11 @@ 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)) - self.assertTrue(job.exists(client=client, retry=retry)) + final_attributes.assert_called_with({'path': "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, None, + job) call_api.assert_called_once_with( retry, @@ -716,8 +729,10 @@ 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) - job.exists(timeout=7.5) + final_attributes.assert_called_with({'path': PATH}, None, job) call_api.assert_called_once_with( DEFAULT_RETRY, @@ -742,8 +757,11 @@ 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() - job.reload() + final_attributes.assert_called_with({'path': "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, None, + job) call_api.assert_called_once_with( DEFAULT_RETRY, @@ -771,8 +789,11 @@ 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) - 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) call_api.assert_called_once_with( retry, @@ -796,8 +817,11 @@ def test_cancel_defaults(self): job = self._set_properties_job() job._properties["jobReference"]["location"] = self.LOCATION connection = job._client._connection = _make_connection(response) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertTrue(job.cancel()) - self.assertTrue(job.cancel()) + final_attributes.assert_called_with({'path': "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)}, + None, job) connection.api_request.assert_called_once_with( method="POST", @@ -821,8 +845,11 @@ def test_cancel_explicit(self): job = self._set_properties_job() client = _make_client(project=other_project) connection = client._connection = _make_connection(response) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertTrue(job.cancel(client=client, timeout=7.5)) - self.assertTrue(job.cancel(client=client, timeout=7.5)) + final_attributes.assert_called_with({'path': "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)}, + None, job) connection.api_request.assert_called_once_with( method="POST", @@ -855,7 +882,11 @@ def test_cancel_w_custom_retry(self): ) with api_request_patcher as fake_api_request: - result = job.cancel(retry=retry, timeout=7.5) + with mock.patch( + 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + result = job.cancel(retry=retry, timeout=7.5) + + final_attributes.assert_called_with({'path': api_path}, None, job) self.assertTrue(result) self.assertEqual(job._properties, resource) @@ -2147,6 +2178,7 @@ def test_result_invokes_begin(self): client._connection = connection job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) + job.result() self.assertEqual(len(connection.api_request.call_args_list), 2) @@ -2343,12 +2375,15 @@ def test_begin_w_bound_client(self): conn = _make_connection(RESOURCE) client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) + path = "/projects/{}/jobs".format(self.PROJECT) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': path}, None, job) conn.api_request.assert_called_once_with( method="POST", - path="/projects/{}/jobs".format(self.PROJECT), + path=path, data={ "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, "configuration": { @@ -2384,7 +2419,10 @@ def test_begin_w_autodetect(self): job = self._make_one( self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client, config ) - job._begin() + 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) sent = { "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, @@ -2478,8 +2516,10 @@ def test_begin_w_alternate_client(self): config.use_avro_logical_types = True config.write_disposition = WriteDisposition.WRITE_TRUNCATE config.schema_update_options = [SchemaUpdateOption.ALLOW_FIELD_ADDITION] + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin(client=client2) - job._begin(client=client2) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() self.assertEqual(len(conn2.api_request.call_args_list), 1) @@ -2504,8 +2544,9 @@ def test_begin_w_job_reference(self): conn = _make_connection(resource) client = _make_client(project=self.PROJECT, connection=conn) load_job = self._make_one(job_ref, [self.SOURCE1], self.TABLE_REF, client) - - load_job._begin() + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + load_job._begin() + final_attributes.assert_called_with({'path': "/projects/alternative-project/jobs"}, None, load_job) conn.api_request.assert_called_once() _, request = conn.api_request.call_args @@ -2522,8 +2563,11 @@ def test_exists_miss_w_bound_client(self): conn = _make_connection() client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertFalse(job.exists()) - self.assertFalse(job.exists()) + final_attributes.assert_called_with({'path': "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, None, + job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None @@ -2536,8 +2580,11 @@ def test_exists_hit_w_alternate_client(self): conn2 = _make_connection({}) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client1) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertTrue(job.exists(client=client2)) - self.assertTrue(job.exists(client=client2)) + final_attributes.assert_called_with({'path': "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, None, + job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -2551,8 +2598,10 @@ def test_exists_miss_w_job_reference(self): conn = _make_connection() client = _make_client(project=self.PROJECT, connection=conn) load_job = self._make_one(job_ref, [self.SOURCE1], self.TABLE_REF, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertFalse(load_job.exists()) - self.assertFalse(load_job.exists()) + final_attributes.assert_called_with({'path': "/projects/other-project/jobs/my-job-id"}, None, load_job) conn.api_request.assert_called_once_with( method="GET", @@ -2567,8 +2616,10 @@ def test_reload_w_bound_client(self): conn = _make_connection(RESOURCE) client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload() - job.reload() + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None @@ -2583,8 +2634,10 @@ def test_reload_w_alternate_client(self): conn2 = _make_connection(RESOURCE) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client1) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload(client=client2) - job.reload(client=client2) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -2602,8 +2655,11 @@ def test_reload_w_job_reference(self): conn = _make_connection(resource) client = _make_client(project=self.PROJECT, connection=conn) load_job = self._make_one(job_ref, [self.SOURCE1], self.TABLE_REF, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + load_job.reload() - load_job.reload() + final_attributes.assert_called_with({'path': "/projects/alternative-project/jobs/{}".format(self.JOB_ID)}, None, + load_job) conn.api_request.assert_called_once_with( method="GET", @@ -2619,8 +2675,10 @@ def test_cancel_w_bound_client(self): conn = _make_connection(RESPONSE) client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.cancel() - job.cancel() + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", path=PATH, query_params={}, timeout=None, @@ -2636,8 +2694,10 @@ def test_cancel_w_alternate_client(self): conn2 = _make_connection(RESPONSE) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client1) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.cancel(client=client2) - job.cancel(client=client2) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -2655,9 +2715,11 @@ def test_cancel_w_job_reference(self): conn = _make_connection({"job": resource}) client = _make_client(project=self.PROJECT, connection=conn) load_job = self._make_one(job_ref, [self.SOURCE1], self.TABLE_REF, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + load_job.cancel() - load_job.cancel() - + final_attributes.assert_called_with( + {'path': "/projects/alternative-project/jobs/{}/cancel".format(self.JOB_ID)}, None, load_job) conn.api_request.assert_called_once_with( method="POST", path="/projects/alternative-project/jobs/{}/cancel".format(self.JOB_ID), @@ -2952,8 +3014,10 @@ def test_begin_w_bound_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -3016,7 +3080,10 @@ def test_begin_w_alternate_client(self): config.create_disposition = CreateDisposition.CREATE_NEVER config.write_disposition = WriteDisposition.WRITE_TRUNCATE job = self._make_one(self.JOB_ID, [source], destination, client1, config) - job._begin(client=client2) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin(client=client2) + + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3038,8 +3105,10 @@ def test_exists_miss_w_bound_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertFalse(job.exists()) - self.assertFalse(job.exists()) + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None, @@ -3054,8 +3123,10 @@ def test_exists_hit_w_alternate_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client1) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertTrue(job.exists(client=client2)) - self.assertTrue(job.exists(client=client2)) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3070,8 +3141,10 @@ def test_reload_w_bound_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload() - job.reload() + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None @@ -3088,8 +3161,10 @@ def test_reload_w_alternate_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client1) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload(client=client2) - job.reload(client=client2) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3349,8 +3424,10 @@ def test_begin_w_bound_client(self): source_dataset = DatasetReference(self.PROJECT, self.DS_ID) source = source_dataset.table(self.SOURCE_TABLE) job = self._make_one(self.JOB_ID, source, [self.DESTINATION_URI], client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -3407,8 +3484,10 @@ def test_begin_w_alternate_client(self): job = self._make_one( self.JOB_ID, source, [self.DESTINATION_URI], client1, config ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin(client=client2) - job._begin(client=client2) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3429,8 +3508,10 @@ def test_exists_miss_w_bound_client(self): job = self._make_one( self.JOB_ID, self.TABLE_REF, [self.DESTINATION_URI], client ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertFalse(job.exists()) - self.assertFalse(job.exists()) + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None, @@ -3445,8 +3526,10 @@ def test_exists_hit_w_alternate_client(self): job = self._make_one( self.JOB_ID, self.TABLE_REF, [self.DESTINATION_URI], client1 ) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertTrue(job.exists(client=client2)) - self.assertTrue(job.exists(client=client2)) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3463,9 +3546,10 @@ def test_reload_w_bound_client(self): source_dataset = DatasetReference(self.PROJECT, self.DS_ID) source = source_dataset.table(self.SOURCE_TABLE) job = self._make_one(self.JOB_ID, source, [self.DESTINATION_URI], client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload() - job.reload() - + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None ) @@ -3483,8 +3567,10 @@ def test_reload_w_alternate_client(self): source_dataset = DatasetReference(self.PROJECT, self.DS_ID) source = source_dataset.table(self.SOURCE_TABLE) job = self._make_one(self.JOB_ID, source, [self.DESTINATION_URI], client1) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload(client=client2) - job.reload(client=client2) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -4023,7 +4109,7 @@ def test_done_w_timeout(self): job = self._get_target_class().from_api_repr(resource, client) with mock.patch.object( - client, "_get_query_results" + client, "_get_query_results" ) as fake_get_results, mock.patch.object(job, "reload") as fake_reload: job.done(timeout=42) @@ -4041,7 +4127,7 @@ def test_done_w_timeout_and_longer_internal_api_timeout(self): job._done_timeout = 8.8 with mock.patch.object( - client, "_get_query_results" + client, "_get_query_results" ) as fake_get_results, mock.patch.object(job, "reload") as fake_reload: job.done(timeout=5.5) @@ -4823,8 +4909,10 @@ def test__begin_w_timeout(self): conn = _make_connection(RESOURCE) client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, self.QUERY, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin(timeout=7.5) - job._begin(timeout=7.5) + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -4856,8 +4944,10 @@ def test_begin_w_bound_client(self): config = QueryJobConfig() config.default_dataset = DatasetReference(self.PROJECT, DS_ID) job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=config) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': PATH}, None, job) self.assertIsNone(job.default_dataset) self.assertEqual(job.udf_resources, []) @@ -4936,8 +5026,10 @@ def test_begin_w_alternate_client(self): config.maximum_bytes_billed = 123456 config.schema_update_options = [SchemaUpdateOption.ALLOW_FIELD_RELAXATION] job = self._make_one(self.JOB_ID, self.QUERY, client1, job_config=config) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin(client=client2) - job._begin(client=client2) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -4978,8 +5070,10 @@ def test_begin_w_udf(self): config.udf_resources = udf_resources config.use_legacy_sql = True job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=config) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': PATH}, None, job) self.assertEqual(job.udf_resources, udf_resources) conn.api_request.assert_called_once_with( @@ -5028,8 +5122,10 @@ def test_begin_w_named_query_parameter(self): jconfig = QueryJobConfig() jconfig.query_parameters = query_parameters job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=jconfig) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': PATH}, None, job) self.assertEqual(job.query_parameters, query_parameters) conn.api_request.assert_called_once_with( @@ -5072,8 +5168,10 @@ def test_begin_w_positional_query_parameter(self): jconfig = QueryJobConfig() jconfig.query_parameters = query_parameters job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=jconfig) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': PATH}, None, job) self.assertEqual(job.query_parameters, query_parameters) conn.api_request.assert_called_once_with( @@ -5148,8 +5246,10 @@ def test_begin_w_table_defs(self): config.table_definitions = {bt_table: bt_config, csv_table: csv_config} config.use_legacy_sql = True job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=config) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -5187,8 +5287,10 @@ def test_dry_run_query(self): config = QueryJobConfig() config.dry_run = True job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=config) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job._begin() - job._begin() + final_attributes.assert_called_with({'path': PATH}, None, job) self.assertEqual(job.udf_resources, []) conn.api_request.assert_called_once_with( method="POST", @@ -5209,8 +5311,10 @@ def test_exists_miss_w_bound_client(self): conn = _make_connection() client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, self.QUERY, client) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertFalse(job.exists()) - self.assertFalse(job.exists()) + final_attributes.assert_called_with({'path': PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None @@ -5223,8 +5327,10 @@ def test_exists_hit_w_alternate_client(self): conn2 = _make_connection({}) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, self.QUERY, client1) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + self.assertTrue(job.exists(client=client2)) - self.assertTrue(job.exists(client=client2)) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5246,8 +5352,10 @@ def test_reload_w_bound_client(self): config = QueryJobConfig() config.destination = table_ref job = self._make_one(self.JOB_ID, None, client, job_config=config) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload() - job.reload() + final_attributes.assert_called_with({'path': PATH}, None, job) self.assertNotEqual(job.destination, table_ref) @@ -5272,8 +5380,10 @@ def test_reload_w_alternate_client(self): conn2 = _make_connection(RESOURCE) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, self.QUERY, client1) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload(client=client2) - job.reload(client=client2) + final_attributes.assert_called_with({'path': PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5296,8 +5406,10 @@ def test_reload_w_timeout(self): config = QueryJobConfig() config.destination = table_ref job = self._make_one(self.JOB_ID, None, client, job_config=config) + with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + job.reload(timeout=4.2) - job.reload(timeout=4.2) + final_attributes.assert_called_with({'path': PATH}, None, job) self.assertNotEqual(job.destination, table_ref) @@ -6018,17 +6130,17 @@ def test_from_api_repr_normal(self): @pytest.mark.parametrize( "query,expected", ( - (None, False), - ("", False), - ("select name, age from table", False), - ("select name, age from table LIMIT 10;", False), - ("select name, age from table order by other_column;", True), - ("Select name, age From table Order By other_column", True), - ("SELECT name, age FROM table ORDER BY other_column;", True), - ("select name, age from table order\nby other_column", True), - ("Select name, age From table Order\nBy other_column;", True), - ("SELECT name, age FROM table ORDER\nBY other_column", True), - ("SelecT name, age froM table OrdeR \n\t BY other_column;", True), + (None, False), + ("", False), + ("select name, age from table", False), + ("select name, age from table LIMIT 10;", False), + ("select name, age from table order by other_column;", True), + ("Select name, age From table Order By other_column", True), + ("SELECT name, age FROM table ORDER BY other_column;", True), + ("select name, age from table order\nby other_column", True), + ("Select name, age From table Order\nBy other_column;", True), + ("SELECT name, age FROM table ORDER\nBY other_column", True), + ("SelecT name, age froM table OrdeR \n\t BY other_column;", True), ), ) def test__contains_order_by(query, expected): @@ -6047,13 +6159,13 @@ def test__contains_order_by(query, expected): @pytest.mark.parametrize( "query", ( - "select name, age from table order by other_column;", - "Select name, age From table Order By other_column;", - "SELECT name, age FROM table ORDER BY other_column;", - "select name, age from table order\nby other_column;", - "Select name, age From table Order\nBy other_column;", - "SELECT name, age FROM table ORDER\nBY other_column;", - "SelecT name, age froM table OrdeR \n\t BY other_column;", + "select name, age from table order by other_column;", + "Select name, age From table Order By other_column;", + "SELECT name, age FROM table ORDER BY other_column;", + "select name, age from table order\nby other_column;", + "Select name, age From table Order\nBy other_column;", + "SELECT name, age FROM table ORDER\nBY other_column;", + "SelecT name, age froM table OrdeR \n\t BY other_column;", ), ) def test_to_dataframe_bqstorage_preserve_order(query): From d4de9269333c9ea23f6f28da59a1f75d93345613 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Sat, 15 Aug 2020 11:52:16 -0400 Subject: [PATCH 20/41] fixing linting issues --- .../cloud/bigquery/opentelemetry_tracing.py | 4 +- tests/unit/test_client.py | 737 ++++++++++++------ tests/unit/test_job.py | 417 ++++++---- 3 files changed, 752 insertions(+), 406 deletions(-) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 530cbed3c..44ed831bc 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -50,9 +50,7 @@ def create_span(name, attributes=None, client=None, job_ref=None): return final_attributes = _get_final_span_attributes(attributes, client, job_ref) - with tracer.start_as_current_span( - name=name, attributes=final_attributes - ) as span: + with tracer.start_as_current_span(name=name, attributes=final_attributes) as span: try: yield span span.set_status(Status(http_status_to_canonical_code(200))) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 88b49fa08..4166e76df 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -57,7 +57,6 @@ bigquery_storage_v1 = None from test_utils.imports import maybe_fail_import from tests.unit.helpers import make_connection -import pdb PANDAS_MINIUM_VERSION = pkg_resources.parse_version("1.0.0") PANDAS_INSTALLED_VERSION = pkg_resources.get_distribution("pandas").parsed_version @@ -249,7 +248,8 @@ def test__get_query_results_miss_w_explicit_project_and_timeout(self): path = "/projects/other-project/queries/nothere" with self.assertRaises(NotFound): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client._get_query_results( "nothere", None, @@ -258,7 +258,7 @@ def test__get_query_results_miss_w_explicit_project_and_timeout(self): timeout_ms=500, timeout=42, ) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_once_with( method="GET", @@ -318,9 +318,11 @@ def test_get_service_account_email(self): email = "bq-123@bigquery-encryption.iam.gserviceaccount.com" resource = {"kind": "bigquery#getServiceAccountResponse", "email": email} conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: service_account_email = client.get_service_account_email(timeout=7.5) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_once_with(method="GET", path=path, timeout=7.5) self.assertEqual(service_account_email, email) @@ -333,9 +335,11 @@ def test_get_service_account_email_w_alternate_project(self): email = "bq-123@bigquery-encryption.iam.gserviceaccount.com" resource = {"kind": "bigquery#getServiceAccountResponse", "email": email} conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: service_account_email = client.get_service_account_email(project=project) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_once_with(method="GET", path=path, timeout=None) self.assertEqual(service_account_email, email) @@ -361,11 +365,12 @@ def test_get_service_account_email_w_custom_retry(self): with api_request_patcher as fake_api_request: with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: service_account_email = client.get_service_account_email( retry=retry, timeout=7.5 ) - final_attributes.assert_called_once_with({'path': api_path}, client, None) + final_attributes.assert_called_once_with({"path": api_path}, client, None) self.assertEqual( service_account_email, "bq-123@bigquery-encryption.iam.gserviceaccount.com" ) @@ -405,9 +410,11 @@ def test_list_projects_defaults(self): creds = _make_credentials() client = self._make_one(PROJECT_1, creds) conn = client._connection = make_connection(DATA) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_projects() - final_attributes.assert_called_once_with({'path': '/projects'}, client, None) + final_attributes.assert_called_once_with({"path": "/projects"}, client, None) page = six.next(iterator.pages) projects = list(page) token = iterator.next_page_token @@ -434,9 +441,11 @@ def test_list_projects_w_timeout(self): creds = _make_credentials() client = self._make_one(PROJECT_1, creds) conn = client._connection = make_connection(DATA) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_projects(timeout=7.5) - final_attributes.assert_called_once_with({'path': '/projects'}, client, None) + final_attributes.assert_called_once_with({"path": "/projects"}, client, None) six.next(iterator.pages) @@ -451,9 +460,11 @@ def test_list_projects_explicit_response_missing_projects_key(self): client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_projects(max_results=3, page_token=TOKEN) - final_attributes.assert_called_once_with({'path': '/projects'}, client, None) + final_attributes.assert_called_once_with({"path": "/projects"}, client, None) page = six.next(iterator.pages) projects = list(page) @@ -502,9 +513,13 @@ def test_list_datasets_defaults(self): creds = _make_credentials() client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_datasets() - final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': None}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % PATH, "page_token": None}, client, None + ) page = six.next(iterator.pages) datasets = list(page) @@ -526,7 +541,9 @@ def test_list_datasets_w_project_and_timeout(self): client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection({}) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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() @@ -547,11 +564,15 @@ def test_list_datasets_explicit_response_missing_datasets_key(self): client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_datasets( include_all=True, filter=FILTER, max_results=3, page_token=TOKEN ) - final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': TOKEN}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % PATH, "page_token": TOKEN}, client, None + ) page = six.next(iterator.pages) datasets = list(page) @@ -587,8 +608,8 @@ def test_dataset_with_specified_project(self): warning for warning in warned if warning.category in (DeprecationWarning, PendingDeprecationWarning) - and "Client.dataset" in str(warning) - and "my_project.my_dataset" in str(warning) + and "Client.dataset" in str(warning) + and "my_project.my_dataset" in str(warning) ] assert matches, "A Client.dataset deprecation warning was not raised." self.assertIsInstance(dataset, DatasetReference) @@ -610,8 +631,8 @@ def test_dataset_with_default_project(self): warning for warning in warned if warning.category in (DeprecationWarning, PendingDeprecationWarning) - and "Client.dataset" in str(warning) - and "my_project.my_dataset" in str(warning) + and "Client.dataset" in str(warning) + and "my_project.my_dataset" in str(warning) ] assert matches, "A Client.dataset deprecation warning was not raised." self.assertIsInstance(dataset, DatasetReference) @@ -631,10 +652,12 @@ def test_get_dataset(self): } conn = client._connection = make_connection(resource) dataset_ref = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: dataset = client.get_dataset(dataset_ref, timeout=7.5) - final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 @@ -647,17 +670,23 @@ def test_get_dataset(self): client._connection = make_connection(Exception(""), resource) with self.assertRaises(Exception): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path}, client, None + ) # Zero-length errors field. client._connection = make_connection(ServerError(""), resource) with self.assertRaises(ServerError): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with({'path': '/%s' % path}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path}, client, None + ) # Non-retryable reason. client._connection = make_connection( @@ -665,10 +694,11 @@ def test_get_dataset(self): ) with self.assertRaises(ServerError): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) # Retryable reason, but retry is disabled. client._connection = make_connection( @@ -676,21 +706,26 @@ def test_get_dataset(self): ) with self.assertRaises(ServerError): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.get_dataset(dataset_ref, retry=None) - final_attributes.assert_called_once_with({'path': '/%s' % path}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path}, client, None + ) # Retryable reason, default retry: success. client._connection = make_connection( ServerError("", errors=[{"reason": "backendError"}]), resource ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: dataset = client.get_dataset( # Test with a string for dataset ID. dataset_ref.dataset_id ) - final_attributes.assert_called_once_with({'path': '/%s' % path}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) @@ -705,7 +740,7 @@ def test_create_bqstorage_client(self): client = self._make_one(project=self.PROJECT, credentials=creds) with mock.patch( - "google.cloud.bigquery_storage_v1.BigQueryReadClient", mock_client + "google.cloud.bigquery_storage_v1.BigQueryReadClient", mock_client ): bqstorage_client = client._create_bqstorage_client() @@ -719,7 +754,7 @@ def test_create_bqstorage_client_missing_dependency(self): def fail_bqstorage_import(name, globals, locals, fromlist, level): # NOTE: *very* simplified, assuming a straightforward absolute import return "bigquery_storage_v1" in name or ( - fromlist is not None and "bigquery_storage_v1" in fromlist + fromlist is not None and "bigquery_storage_v1" in fromlist ) no_bqstorage = maybe_fail_import(predicate=fail_bqstorage_import) @@ -732,7 +767,7 @@ def fail_bqstorage_import(name, globals, locals, fromlist, level): warning for warning in warned if "not installed" in str(warning) - and "google-cloud-bigquery-storage" in str(warning) + and "google-cloud-bigquery-storage" in str(warning) ] assert matching_warnings, "Missing dependency warning not raised." @@ -751,10 +786,12 @@ def test_create_dataset_minimal(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) before = Dataset(ds_ref) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: after = client.create_dataset(before, timeout=7.5) - final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -815,9 +852,11 @@ def test_create_dataset_w_attrs(self): before.default_table_expiration_ms = 3600 before.location = LOCATION before.labels = LABELS - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: after = client.create_dataset(before) - final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -867,10 +906,12 @@ def test_create_dataset_w_custom_property(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) before = Dataset(ds_ref) before._properties["newAlphaProperty"] = "unreleased property" - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: after = client.create_dataset(before) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -908,10 +949,12 @@ def test_create_dataset_w_client_location_wo_dataset_location(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) before = Dataset(ds_ref) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: after = client.create_dataset(before) - final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -953,9 +996,11 @@ def test_create_dataset_w_client_location_w_dataset_location(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) before = Dataset(ds_ref) before.location = OTHER_LOCATION - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: after = client.create_dataset(before) - final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) self.assertEqual(after.project, self.PROJECT) @@ -990,10 +1035,12 @@ def test_create_dataset_w_reference(self): project=self.PROJECT, credentials=creds, location=self.LOCATION ) conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: dataset = client.create_dataset(DatasetReference(self.PROJECT, self.DS_ID)) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -1028,10 +1075,12 @@ def test_create_dataset_w_fully_qualified_string(self): project=self.PROJECT, credentials=creds, location=self.LOCATION ) conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: dataset = client.create_dataset("{}.{}".format(self.PROJECT, self.DS_ID)) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -1066,10 +1115,12 @@ def test_create_dataset_w_string(self): project=self.PROJECT, credentials=creds, location=self.LOCATION ) conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: dataset = client.create_dataset(self.DS_ID) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -1102,10 +1153,11 @@ def test_create_dataset_alreadyexists_w_exists_ok_false(self): with pytest.raises(google.api_core.exceptions.AlreadyExists): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.create_dataset(self.DS_ID) - final_attributes.assert_called_once_with({'path': 'path'}, client, None) + final_attributes.assert_called_once_with({"path": "path"}, client, None) def test_create_dataset_alreadyexists_w_exists_ok_true(self): post_path = "/projects/{}/datasets".format(self.PROJECT) @@ -1123,9 +1175,11 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(self): conn = client._connection = make_connection( google.api_core.exceptions.AlreadyExists("dataset already exists"), resource ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: dataset = client.create_dataset(self.DS_ID, exists_ok=True) - final_attributes.assert_called_with({'path': get_path}, client, None) + final_attributes.assert_called_with({"path": get_path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) self.assertEqual(dataset.project, self.PROJECT) @@ -1169,16 +1223,17 @@ def test_create_routine_w_minimal_resource(self): conn = client._connection = make_connection(resource) full_routine_id = "test-routine-project.test_routines.minimal_routine" routine = Routine(full_routine_id) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: actual_routine = client.create_routine(routine, timeout=7.5) - final_attributes.assert_called_once_with({'path': path, 'exists_ok': False}, client, None) + final_attributes.assert_called_once_with( + {"path": path, "exists_ok": False}, client, None + ) conn.api_request.assert_called_once_with( - method="POST", - path=path, - data=resource, - timeout=7.5, + method="POST", path=path, data=resource, timeout=7.5, ) self.assertEqual( actual_routine.reference, RoutineReference.from_string(full_routine_id) @@ -1192,16 +1247,19 @@ def test_create_routine_w_conflict(self): conn = client._connection = make_connection( google.api_core.exceptions.AlreadyExists("routine already exists") ) - path = '/projects/test-routine-project/datasets/test_routines/routines' + path = "/projects/test-routine-project/datasets/test_routines/routines" full_routine_id = "test-routine-project.test_routines.minimal_routine" routine = Routine(full_routine_id) with pytest.raises(google.api_core.exceptions.AlreadyExists): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.create_routine(routine) - final_attributes.assert_called_once_with({'path': path, 'exists_okay': False}, client, None) + final_attributes.assert_called_once_with( + {"path": path, "exists_okay": False}, client, None + ) resource = { "routineReference": { @@ -1211,10 +1269,7 @@ def test_create_routine_w_conflict(self): } } conn.api_request.assert_called_once_with( - method="POST", - path=path, - data=resource, - timeout=None, + method="POST", path=path, data=resource, timeout=None, ) def test_create_routine_w_conflict_exists_ok(self): @@ -1236,22 +1291,21 @@ def test_create_routine_w_conflict_exists_ok(self): ) full_routine_id = "test-routine-project.test_routines.minimal_routine" routine = Routine(full_routine_id) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: actual_routine = client.create_routine(routine, exists_ok=True) - final_attributes.assert_called_with({'path': path + '/minimal_routine'}, client, None) + final_attributes.assert_called_with( + {"path": path + "/minimal_routine"}, client, None + ) self.assertEqual(actual_routine.project, "test-routine-project") self.assertEqual(actual_routine.dataset_id, "test_routines") self.assertEqual(actual_routine.routine_id, "minimal_routine") conn.api_request.assert_has_calls( [ - mock.call( - method="POST", - path=path, - data=resource, - timeout=None, - ), + mock.call(method="POST", path=path, data=resource, timeout=None,), mock.call( method="GET", path="/projects/test-routine-project/datasets/test_routines/routines/minimal_routine", @@ -1271,10 +1325,14 @@ def test_create_table_w_day_partition(self): conn = client._connection = make_connection(resource) table = Table(self.TABLE_REF) table.time_partitioning = TimePartitioning() - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table(table, timeout=7.5) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None + ) conn.api_request.assert_called_once_with( method="POST", @@ -1306,9 +1364,13 @@ def test_create_table_w_custom_property(self): conn = client._connection = make_connection(resource) table = Table(self.TABLE_REF) table._properties["newAlphaProperty"] = "unreleased property" - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None + ) conn.api_request.assert_called_once_with( method="POST", @@ -1342,10 +1404,14 @@ def test_create_table_w_encryption_configuration(self): table.encryption_configuration = EncryptionConfiguration( kms_key_name=self.KMS_KEY_NAME ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None + ) conn.api_request.assert_called_once_with( method="POST", @@ -1374,10 +1440,14 @@ def test_create_table_w_day_partition_and_expire(self): conn = client._connection = make_connection(resource) table = Table(self.TABLE_REF) table.time_partitioning = TimePartitioning(expiration_ms=100) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None + ) conn.api_request.assert_called_once_with( method="POST", @@ -1435,10 +1505,14 @@ def test_create_table_w_schema_and_query(self): table = Table(self.TABLE_REF, schema=schema) table.view_query = query - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None + ) conn.api_request.assert_called_once_with( method="POST", @@ -1499,10 +1573,14 @@ def test_create_table_w_external(self): ec.autodetect = True table.external_data_configuration = ec - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table(table) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': table.dataset_id}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None + ) conn.api_request.assert_called_once_with( method="POST", @@ -1536,11 +1614,16 @@ def test_create_table_w_reference(self): resource = self._make_table_resource() conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table(self.TABLE_REF) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': self.TABLE_REF.dataset_id}, - client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": self.TABLE_REF.dataset_id}, + client, + None, + ) conn.api_request.assert_called_once_with( method="POST", @@ -1563,12 +1646,17 @@ def test_create_table_w_fully_qualified_string(self): client = self._make_one(project=self.PROJECT, credentials=creds) resource = self._make_table_resource() conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table( "{}.{}.{}".format(self.PROJECT, self.DS_ID, self.TABLE_ID) ) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': self.TABLE_REF.dataset_id}, - client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": self.TABLE_REF.dataset_id}, + client, + None, + ) conn.api_request.assert_called_once_with( method="POST", @@ -1591,11 +1679,16 @@ def test_create_table_w_string(self): client = self._make_one(project=self.PROJECT, credentials=creds) resource = self._make_table_resource() conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'dataset_id': self.TABLE_REF.dataset_id}, - client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "dataset_id": self.TABLE_REF.dataset_id}, + client, + None, + ) conn.api_request.assert_called_once_with( method="POST", @@ -1624,11 +1717,15 @@ def test_create_table_alreadyexists_w_exists_ok_false(self): with pytest.raises(google.api_core.exceptions.AlreadyExists): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_with({'path': post_path, 'dataset_id': self.TABLE_REF.dataset_id}, client, - None) + final_attributes.assert_called_with( + {"path": post_path, "dataset_id": self.TABLE_REF.dataset_id}, + client, + None, + ) conn.api_request.assert_called_once_with( method="POST", @@ -1658,11 +1755,13 @@ def test_create_table_alreadyexists_w_exists_ok_true(self): google.api_core.exceptions.AlreadyExists("table already exists"), resource ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.create_table( "{}.{}".format(self.DS_ID, self.TABLE_ID), exists_ok=True ) - final_attributes.assert_called_with({'path': get_path}, client, None) + final_attributes.assert_called_with({"path": get_path}, client, None) self.assertEqual(got.project, self.PROJECT) self.assertEqual(got.dataset_id, self.DS_ID) @@ -1717,10 +1816,12 @@ def test_get_model(self): conn = client._connection = make_connection(resource) model_ref = DatasetReference(self.PROJECT, self.DS_ID).model(self.MODEL_ID) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.get_model(model_ref, timeout=7.5) - final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 @@ -1746,10 +1847,12 @@ def test_get_model_w_string(self): conn = client._connection = make_connection(resource) model_id = "{}.{}.{}".format(self.PROJECT, self.DS_ID, self.MODEL_ID) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: got = client.get_model(model_id) - final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=None @@ -1783,15 +1886,14 @@ def test_get_routine(self): conn = client._connection = make_connection(resource) with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: actual_routine = client.get_routine(routine, timeout=7.5) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_once_with( - method="GET", - path=path, - timeout=7.5, + method="GET", path=path, timeout=7.5, ) self.assertEqual( actual_routine.reference, @@ -1820,10 +1922,12 @@ def test_get_table(self): client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) resource = self._make_table_resource() conn = client._connection = make_connection(resource) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: table = client.get_table(self.TABLE_REF, timeout=7.5) - final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) conn.api_request.assert_called_once_with( method="GET", path="/%s" % path, timeout=7.5 @@ -1898,10 +2002,14 @@ def test_get_iam_policy(self): http = object() client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(RETURNED) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: policy = client.get_iam_policy(self.TABLE_REF, timeout=7.5) - final_attributes.assert_called_once_with({'path': PATH, 'body': BODY}, client, None) + final_attributes.assert_called_once_with( + {"path": PATH, "body": BODY}, client, None + ) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -1970,11 +2078,13 @@ def test_set_iam_policy(self): client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(RETURNED) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: returned_policy = client.set_iam_policy( self.TABLE_REF, policy, updateMask=MASK, timeout=7.5 ) - final_attributes.assert_called_once_with({'path': PATH}, client, None) + final_attributes.assert_called_once_with({"path": PATH}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -2000,10 +2110,12 @@ def test_set_iam_policy_no_mask(self): http = object() client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(RETURNED) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.set_iam_policy(self.TABLE_REF, policy, timeout=7.5) - final_attributes.assert_called_once_with({'path': PATH}, client, None) + final_attributes.assert_called_once_with({"path": PATH}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -2021,11 +2133,13 @@ def test_set_iam_policy_invalid_policy(self): with self.assertRaises(TypeError): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.set_iam_policy(self.TABLE_REF, invalid_policy_repr) - final_attributes.assert_called_once_with({'path': "{}:setIamPolicy".format(self.TABLE_REF.path)}, client, - None) + final_attributes.assert_called_once_with( + {"path": "{}:setIamPolicy".format(self.TABLE_REF.path)}, client, None + ) def test_set_iam_policy_w_invalid_table(self): from google.api_core.iam import Policy @@ -2044,10 +2158,13 @@ def test_set_iam_policy_w_invalid_table(self): with self.assertRaises(TypeError): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.set_iam_policy(table_resource_string, policy) - final_attributes.assert_called_once_with({'path': table_resource_string}, client, None) + final_attributes.assert_called_once_with( + {"path": table_resource_string}, client, None + ) def test_test_iam_permissions(self): PATH = "/projects/%s/datasets/%s/tables/%s:testIamPermissions" % ( @@ -2064,10 +2181,14 @@ def test_test_iam_permissions(self): http = object() client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) conn = client._connection = make_connection(RETURNED) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.test_iam_permissions(self.TABLE_REF, PERMISSIONS, timeout=7.5) - final_attributes.assert_called_once_with({'path': PATH, 'body': BODY}, client, None) + final_attributes.assert_called_once_with( + {"path": PATH, "body": BODY}, client, None + ) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -2129,15 +2250,21 @@ def test_update_dataset(self): ds.default_table_expiration_ms = EXP ds.labels = LABELS ds.access_entries = [AccessEntry("OWNER", "userByEmail", "phred@example.com")] - fields = ["description", "friendly_name", "location", "labels", "access_entries"] + fields = [ + "description", + "friendly_name", + "location", + "labels", + "access_entries", + ] - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: - ds2 = client.update_dataset( - ds, - fields=fields, - timeout=7.5, - ) - final_attributes.assert_called_once_with({'path': "/" + PATH, 'fields': fields}, client, None) + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: + ds2 = client.update_dataset(ds, fields=fields, timeout=7.5,) + final_attributes.assert_called_once_with( + {"path": "/" + PATH, "fields": fields}, client, None + ) conn.api_request.assert_called_once_with( method="PATCH", @@ -2180,10 +2307,14 @@ def test_update_dataset_w_custom_property(self): dataset = Dataset(DatasetReference(self.PROJECT, self.DS_ID)) dataset._properties["newAlphaProperty"] = "unreleased property" - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: dataset = client.update_dataset(dataset, ["newAlphaProperty"]) - final_attributes.assert_called_once_with({'path': path, 'fields': ["newAlphaProperty"]}, client, None) + final_attributes.assert_called_once_with( + {"path": path, "fields": ["newAlphaProperty"]}, client, None + ) conn.api_request.assert_called_once_with( method="PATCH", @@ -2232,11 +2363,13 @@ def test_update_model(self): model.expires = expires model.labels = {"x": "y"} fields = ["description", "friendly_name", "labels", "expires"] - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: - updated_model = client.update_model( - model, fields, timeout=7.5 - ) - final_attributes.assert_called_once_with({'path': "/" + path, 'fields': fields}, client, None) + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: + updated_model = client.update_model(model, fields, timeout=7.5) + final_attributes.assert_called_once_with( + {"path": "/" + path, "fields": fields}, client, None + ) sent = { "description": description, @@ -2293,15 +2426,22 @@ def test_update_routine(self): routine.language = "SQL" routine.type_ = "SCALAR_FUNCTION" routine._properties["someNewField"] = "someValue" - fields = ["arguments", "language", "body", "type_", "return_type", "someNewField"] + fields = [ + "arguments", + "language", + "body", + "type_", + "return_type", + "someNewField", + ] - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: - actual_routine = client.update_routine( - routine, - fields, - timeout=7.5, - ) - final_attributes.assert_called_once_with({'path': routine.path, 'fields': fields}, client, None) + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: + actual_routine = client.update_routine(routine, fields, timeout=7.5,) + final_attributes.assert_called_once_with( + {"path": routine.path, "fields": fields}, client, None + ) # TODO: routineReference isn't needed when the Routines API supports # partial updates. @@ -2320,10 +2460,14 @@ def test_update_routine(self): # ETag becomes If-Match header. routine._properties["etag"] = "im-an-etag" - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.update_routine(routine, []) - final_attributes.assert_called_once_with({'path': routine.path, 'fields': []}, client, None) + final_attributes.assert_called_once_with( + {"path": routine.path, "fields": []}, client, None + ) req = conn.api_request.call_args self.assertEqual(req[1]["headers"]["If-Match"], "im-an-etag") @@ -2376,12 +2520,14 @@ def test_update_table(self): table.friendly_name = title table.labels = {"x": "y"} fields = ["schema", "description", "friendly_name", "labels"] - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: - updated_table = client.update_table( - table, fields, timeout=7.5 - ) - span_path = '/%s' % path - final_attributes.assert_called_once_with({'path': span_path, 'fields': fields}, client, None) + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: + updated_table = client.update_table(table, fields, timeout=7.5) + span_path = "/%s" % path + final_attributes.assert_called_once_with( + {"path": span_path, "fields": fields}, client, None + ) sent = { "schema": { @@ -2414,9 +2560,13 @@ def test_update_table(self): # ETag becomes If-Match header. table._properties["etag"] = "etag" - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.update_table(table, []) - final_attributes.assert_called_once_with({'path': "/" + path, 'fields': []}, client, None) + final_attributes.assert_called_once_with( + {"path": "/" + path, "fields": []}, client, None + ) req = conn.api_request.call_args self.assertEqual(req[1]["headers"]["If-Match"], "etag") @@ -2437,10 +2587,14 @@ def test_update_table_w_custom_property(self): table = Table(self.TABLE_REF) table._properties["newAlphaProperty"] = "unreleased property" - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: updated_table = client.update_table(table, ["newAlphaProperty"]) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["newAlphaProperty"]}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "fields": ["newAlphaProperty"]}, client, None + ) conn.api_request.assert_called_once_with( method="PATCH", @@ -2468,11 +2622,14 @@ def test_update_table_only_use_legacy_sql(self): conn = client._connection = make_connection(resource) table = Table(self.TABLE_REF) table.view_use_legacy_sql = True - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: updated_table = client.update_table(table, ["view_use_legacy_sql"]) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["view_use_legacy_sql"]}, client, - None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "fields": ["view_use_legacy_sql"]}, client, None + ) conn.api_request.assert_called_once_with( method="PATCH", @@ -2535,10 +2692,14 @@ def test_update_table_w_query(self): table.view_query = query table.view_use_legacy_sql = True updated_properties = ["schema", "view_query", "expires", "view_use_legacy_sql"] - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: updated_table = client.update_table(table, updated_properties) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': updated_properties}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "fields": updated_properties}, client, None + ) self.assertEqual(updated_table.schema, table.schema) self.assertEqual(updated_table.view_query, table.view_query) @@ -2581,7 +2742,9 @@ def test_update_table_w_schema_None(self): creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(resource1, resource2) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: table = client.get_table( # Test with string for table ID "{}.{}.{}".format( @@ -2590,14 +2753,18 @@ def test_update_table_w_schema_None(self): self.TABLE_REF.table_id, ) ) - final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) table.schema = None - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: updated_table = client.update_table(table, ["schema"]) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["schema"]}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "fields": ["schema"]}, client, None + ) self.assertEqual(len(conn.api_request.call_args_list), 2) req = conn.api_request.call_args_list[1] @@ -2628,19 +2795,28 @@ def test_update_table_delete_property(self): table.description = description table.friendly_name = title - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: table2 = client.update_table(table, ["description", "friendly_name"]) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["description", "friendly_name"]}, - client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "fields": ["description", "friendly_name"]}, + client, + None, + ) self.assertEqual(table2.description, table.description) table2.description = None - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: table3 = client.update_table(table2, ["description"]) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'fields': ["description"]}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "fields": ["description"]}, client, None + ) self.assertEqual(len(conn.api_request.call_args_list), 2) req = conn.api_request.call_args_list[1] @@ -2657,10 +2833,14 @@ def test_list_tables_empty_w_timeout(self): conn = client._connection = make_connection({}) dataset = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_tables(dataset, timeout=7.5) - final_attributes.assert_called_once_with({'path': path, 'page_token': None, 'max_results': None}, client, None) + final_attributes.assert_called_once_with( + {"path": path, "page_token": None, "max_results": None}, client, None + ) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2680,11 +2860,16 @@ def test_list_models_empty_w_timeout(self): conn = client._connection = make_connection({}) dataset_id = "{}.{}".format(self.PROJECT, self.DS_ID) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_models(dataset_id, timeout=7.5) final_attributes.assert_called_once_with( - {'path': path, 'page_token': iterator.next_page_token, 'max_results': None}, client, None) + {"path": path, "page_token": iterator.next_page_token, "max_results": None}, + client, + None, + ) page = six.next(iterator.pages) models = list(page) @@ -2728,11 +2913,16 @@ def test_list_models_defaults(self): conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_models(dataset) - final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': None, 'max_results': None}, - client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % PATH, "page_token": None, "max_results": None}, + client, + None, + ) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2761,10 +2951,14 @@ def test_list_routines_empty_w_timeout(self): conn = client._connection = make_connection({}) path = "/projects/test-routines/datasets/test_routines/routines" - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_routines("test-routines.test_routines", timeout=7.5) - final_attributes.assert_called_once_with({'path': path, 'page_token': None, 'max_results': None}, client, None) + final_attributes.assert_called_once_with( + {"path": path, "page_token": None, "max_results": None}, client, None + ) page = six.next(iterator.pages) routines = list(page) @@ -2773,10 +2967,7 @@ def test_list_routines_empty_w_timeout(self): self.assertEqual(routines, []) self.assertIsNone(token) conn.api_request.assert_called_once_with( - method="GET", - path=path, - query_params={}, - timeout=7.5, + method="GET", path=path, query_params={}, timeout=7.5, ) def test_list_routines_defaults(self): @@ -2812,10 +3003,14 @@ def test_list_routines_defaults(self): client = self._make_one(project=project_id, credentials=creds) conn = client._connection = make_connection(resource) dataset = DatasetReference(client.project, dataset_id) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_routines(dataset) - final_attributes.assert_called_once_with({'path': path, 'page_token': None, 'max_results': None}, client, None) + final_attributes.assert_called_once_with( + {"path": path, "page_token": None, "max_results": None}, client, None + ) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2879,11 +3074,16 @@ def test_list_tables_defaults(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_tables(dataset) - final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': None, 'max_results': None}, - client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % PATH, "page_token": None, "max_results": None}, + client, + None, + ) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2937,15 +3137,18 @@ def test_list_tables_explicit(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: iterator = client.list_tables( # Test with string for dataset ID. self.DS_ID, max_results=3, page_token=TOKEN, ) - final_attributes.assert_called_once_with({'path': "/%s" % PATH, 'page_token': TOKEN, 'max_results': 3}, client, - None) + final_attributes.assert_called_once_with( + {"path": "/%s" % PATH, "page_token": TOKEN, "max_results": 3}, client, None + ) self.assertEqual(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2984,10 +3187,13 @@ def test_delete_dataset(self): conn = client._connection = make_connection(*([{}] * len(datasets))) for arg in datasets: with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_dataset(arg, timeout=7.5) - final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % PATH}, client, None + ) conn.api_request.assert_called_with( method="DELETE", path="/%s" % PATH, query_params={}, timeout=7.5 @@ -3003,10 +3209,13 @@ def test_delete_dataset_delete_contents(self): ds_ref = DatasetReference(self.PROJECT, self.DS_ID) for arg in (ds_ref, Dataset(ds_ref)): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_dataset(arg, delete_contents=True) - final_attributes.assert_called_once_with({'path': "/%s" % PATH}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % PATH}, client, None + ) conn.api_request.assert_called_with( method="DELETE", path="/%s" % PATH, @@ -3033,10 +3242,11 @@ def test_delete_dataset_w_not_found_ok_false(self): with self.assertRaises(google.api_core.exceptions.NotFound): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_dataset(self.DS_ID) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with( method="DELETE", path=path, query_params={}, timeout=None @@ -3051,10 +3261,12 @@ def test_delete_dataset_w_not_found_ok_true(self): google.api_core.exceptions.NotFound("dataset not found") ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_dataset(self.DS_ID, not_found_ok=True) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with( method="DELETE", path=path, query_params={}, timeout=None @@ -3081,10 +3293,13 @@ def test_delete_model(self): for arg in models: with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_model(arg, timeout=7.5) - final_attributes.assert_called_once_with({'path': "/%s" % path, 'not_found_okay': False}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path, "not_found_okay": False}, client, None + ) conn.api_request.assert_called_with( method="DELETE", path="/%s" % path, timeout=7.5 ) @@ -3121,11 +3336,15 @@ def test_delete_model_w_not_found_ok_true(self): conn = client._connection = make_connection( google.api_core.exceptions.NotFound("model not found") ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_model( "{}.{}".format(self.DS_ID, self.MODEL_ID), not_found_ok=True ) - final_attributes.assert_called_once_with({'path': path, 'not_found_okay': True}, client, None) + final_attributes.assert_called_once_with( + {"path": path, "not_found_okay": True}, client, None + ) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) @@ -3147,15 +3366,14 @@ def test_delete_routine(self): for routine in routines: with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_routine(routine, timeout=7.5) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with( - method="DELETE", - path=path, - timeout=7.5, + method="DELETE", path=path, timeout=7.5, ) def test_delete_routine_w_wrong_type(self): @@ -3175,15 +3393,14 @@ def test_delete_routine_w_not_found_ok_false(self): with self.assertRaises(google.api_core.exceptions.NotFound): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_routine("routines-project.test_routines.test_routine") - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with( - method="DELETE", - path=path, - timeout=None, + method="DELETE", path=path, timeout=None, ) def test_delete_routine_w_not_found_ok_true(self): @@ -3195,16 +3412,16 @@ def test_delete_routine_w_not_found_ok_true(self): ) path = "/projects/routines-project/datasets/test_routines/routines/test_routine" - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_routine( "routines-project.test_routines.test_routine", not_found_ok=True ) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with( - method="DELETE", - path=path, - timeout=None, + method="DELETE", path=path, timeout=None, ) def test_delete_table(self): @@ -3231,10 +3448,13 @@ def test_delete_table(self): for arg in tables: with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_table(arg, timeout=7.5) - final_attributes.assert_called_once_with({'path': "/%s" % path}, client, None) + final_attributes.assert_called_once_with( + {"path": "/%s" % path}, client, None + ) conn.api_request.assert_called_with( method="DELETE", path="/%s" % path, timeout=7.5 @@ -3259,10 +3479,11 @@ def test_delete_table_w_not_found_ok_false(self): with self.assertRaises(google.api_core.exceptions.NotFound): with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) @@ -3277,12 +3498,14 @@ def test_delete_table_w_not_found_ok_true(self): google.api_core.exceptions.NotFound("table not found") ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: client.delete_table( "{}.{}".format(self.DS_ID, self.TABLE_ID), not_found_ok=True ) - final_attributes.assert_called_once_with({'path': path}, client, None) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) @@ -4129,9 +4352,9 @@ def _initiate_resumable_upload_helper(self, num_retries=None): # Check the returned values. self.assertIsInstance(upload, ResumableUpload) upload_url = ( - "https://bigquery.googleapis.com/upload/bigquery/v2/projects/" - + self.PROJECT - + "/jobs?uploadType=resumable" + "https://bigquery.googleapis.com/upload/bigquery/v2/projects/" + + self.PROJECT + + "/jobs?uploadType=resumable" ) self.assertEqual(upload.upload_url, upload_url) expected_headers = _get_upload_headers(conn.user_agent) @@ -4198,20 +4421,20 @@ def _do_multipart_upload_success_helper(self, get_boundary, num_retries=None): get_boundary.assert_called_once_with() upload_url = ( - "https://bigquery.googleapis.com/upload/bigquery/v2/projects/" - + self.PROJECT - + "/jobs?uploadType=multipart" + "https://bigquery.googleapis.com/upload/bigquery/v2/projects/" + + self.PROJECT + + "/jobs?uploadType=multipart" ) payload = ( - b"--==0==\r\n" - + b"content-type: application/json; charset=UTF-8\r\n\r\n" - + json.dumps(metadata).encode("utf-8") - + b"\r\n" - + b"--==0==\r\n" - + b"content-type: */*\r\n\r\n" - + data - + b"\r\n" - + b"--==0==--" + b"--==0==\r\n" + + b"content-type: application/json; charset=UTF-8\r\n\r\n" + + json.dumps(metadata).encode("utf-8") + + b"\r\n" + + b"--==0==\r\n" + + b"content-type: */*\r\n\r\n" + + data + + b"\r\n" + + b"--==0==--" ) headers = _get_upload_headers(conn.user_agent) headers["content-type"] = b'multipart/related; boundary="==0=="' @@ -5023,7 +5246,7 @@ def test_query_preserving_explicit_job_config(self): from google.cloud.bigquery import QueryJobConfig - client = self._make_one(project=self.PROJECT, credentials=creds, _http=http, ) + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http,) conn = client._connection = make_connection(resource) job_config = QueryJobConfig() @@ -6040,7 +6263,7 @@ def test_insert_rows_from_dataframe(self): actual_calls = conn.api_request.call_args_list for call, expected_data in six.moves.zip_longest( - actual_calls, EXPECTED_SENT_DATA + actual_calls, EXPECTED_SENT_DATA ): expected_call = mock.call( method="POST", path=API_PATH, data=expected_data, timeout=7.5 @@ -6108,7 +6331,7 @@ def test_insert_rows_from_dataframe_nan(self): actual_calls = conn.api_request.call_args_list for call, expected_data in six.moves.zip_longest( - actual_calls, EXPECTED_SENT_DATA + actual_calls, EXPECTED_SENT_DATA ): expected_call = mock.call( method="POST", path=API_PATH, data=expected_data, timeout=7.5 @@ -7069,7 +7292,7 @@ def test_load_table_from_file_with_rewind(self): file_obj.seek(2) with self._make_do_upload_patch( - client, "_do_resumable_upload", self.EXPECTED_CONFIGURATION + client, "_do_resumable_upload", self.EXPECTED_CONFIGURATION ): client.load_table_from_file(file_obj, self.TABLE_REF, rewind=True) @@ -7527,8 +7750,8 @@ def test_load_table_from_dataframe_no_schema_warning_wo_pyarrow(self): warning for warning in warned if warning.category in (DeprecationWarning, PendingDeprecationWarning) - and "could not be detected" in str(warning) - and "please provide a schema" in str(warning) + and "could not be detected" in str(warning) + and "please provide a schema" in str(warning) ] assert matches, "A missing schema deprecation warning was not raised." @@ -7804,7 +8027,7 @@ def test_load_table_from_dataframe_w_partial_schema_extra_types(self): ) job_config = job.LoadJobConfig(schema=schema) with load_patch as load_table_from_file, pytest.raises( - ValueError + ValueError ) as exc_context: client.load_table_from_dataframe( dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION @@ -7840,7 +8063,7 @@ def test_load_table_from_dataframe_w_partial_schema_missing_types(self): schema = (SchemaField("string_col", "STRING"),) job_config = job.LoadJobConfig(schema=schema) with pyarrow_patch, load_patch as load_table_from_file, warnings.catch_warnings( - record=True + record=True ) as warned: client.load_table_from_dataframe( dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION @@ -7889,7 +8112,7 @@ def test_load_table_from_dataframe_w_schema_wo_pyarrow(self): pyarrow_patch = mock.patch("google.cloud.bigquery.client.pyarrow", None) with load_patch as load_table_from_file, pyarrow_patch, warnings.catch_warnings( - record=True + record=True ) as warned: client.load_table_from_dataframe( dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 4a8bdd467..06b280d9b 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -71,17 +71,17 @@ def _make_connection(*responses): def _make_job_resource( - creation_time_ms=1437767599006, - started_time_ms=1437767600007, - ended_time_ms=1437767601008, - started=False, - ended=False, - etag="abc-def-hjk", - endpoint="https://bigquery.googleapis.com", - job_type="load", - job_id="a-random-id", - project_id="some-project", - user_email="bq-user@example.com", + creation_time_ms=1437767599006, + started_time_ms=1437767600007, + ended_time_ms=1437767601008, + started=False, + ended=False, + etag="abc-def-hjk", + endpoint="https://bigquery.googleapis.com", + job_type="load", + job_id="a-random-id", + project_id="some-project", + user_email="bq-user@example.com", ): resource = { "configuration": {job_type: {}}, @@ -621,17 +621,15 @@ 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: + 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) + final_attributes.assert_called_with({"path": path}, None, job) call_api.assert_called_once_with( - DEFAULT_RETRY, - method="POST", - path=path, - data=resource, - timeout=None, + DEFAULT_RETRY, method="POST", path=path, data=resource, timeout=None, ) self.assertEqual(job._properties, resource) @@ -655,17 +653,15 @@ 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: + 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) + final_attributes.assert_called_with({"path": path}, None, job) call_api.assert_called_once_with( - retry, - method="POST", - path=path, - data=resource, - timeout=7.5, + retry, method="POST", path=path, data=resource, timeout=7.5, ) self.assertEqual(job._properties, resource) @@ -678,11 +674,16 @@ def test_exists_defaults_miss(self): 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: + 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) + final_attributes.assert_called_with( + {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, + None, + job, + ) call_api.assert_called_once_with( DEFAULT_RETRY, @@ -709,11 +710,16 @@ 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: + 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) + final_attributes.assert_called_with( + {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, + None, + job, + ) call_api.assert_called_once_with( retry, @@ -729,10 +735,12 @@ 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: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) call_api.assert_called_once_with( DEFAULT_RETRY, @@ -757,11 +765,16 @@ 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: + 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) + final_attributes.assert_called_with( + {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, + None, + job, + ) call_api.assert_called_once_with( DEFAULT_RETRY, @@ -789,11 +802,16 @@ 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: + 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) + final_attributes.assert_called_with( + {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, + None, + job, + ) call_api.assert_called_once_with( retry, @@ -817,11 +835,16 @@ def test_cancel_defaults(self): job = self._set_properties_job() job._properties["jobReference"]["location"] = self.LOCATION connection = job._client._connection = _make_connection(response) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) 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_with( + {"path": "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)}, + None, + job, + ) connection.api_request.assert_called_once_with( method="POST", @@ -845,11 +868,16 @@ def test_cancel_explicit(self): job = self._set_properties_job() client = _make_client(project=other_project) connection = client._connection = _make_connection(response) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: self.assertTrue(job.cancel(client=client, timeout=7.5)) - final_attributes.assert_called_with({'path': "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)}, - None, job) + final_attributes.assert_called_with( + {"path": "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)}, + None, + job, + ) connection.api_request.assert_called_once_with( method="POST", @@ -883,10 +911,11 @@ def test_cancel_w_custom_retry(self): with api_request_patcher as fake_api_request: with mock.patch( - 'google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) 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_with({"path": api_path}, None, job) self.assertTrue(result) self.assertEqual(job._properties, resource) @@ -2376,10 +2405,12 @@ def test_begin_w_bound_client(self): client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) path = "/projects/{}/jobs".format(self.PROJECT) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": path}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -2419,10 +2450,12 @@ def test_begin_w_autodetect(self): job = self._make_one( self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client, config ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": path}, None, job) sent = { "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, @@ -2516,10 +2549,12 @@ def test_begin_w_alternate_client(self): config.use_avro_logical_types = True config.write_disposition = WriteDisposition.WRITE_TRUNCATE config.schema_update_options = [SchemaUpdateOption.ALLOW_FIELD_ADDITION] - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job._begin(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() self.assertEqual(len(conn2.api_request.call_args_list), 1) @@ -2544,9 +2579,13 @@ def test_begin_w_job_reference(self): conn = _make_connection(resource) client = _make_client(project=self.PROJECT, connection=conn) load_job = self._make_one(job_ref, [self.SOURCE1], self.TABLE_REF, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: load_job._begin() - final_attributes.assert_called_with({'path': "/projects/alternative-project/jobs"}, None, load_job) + final_attributes.assert_called_with( + {"path": "/projects/alternative-project/jobs"}, None, load_job + ) conn.api_request.assert_called_once() _, request = conn.api_request.call_args @@ -2563,11 +2602,16 @@ def test_exists_miss_w_bound_client(self): conn = _make_connection() client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with( + {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, + None, + job, + ) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None @@ -2580,11 +2624,16 @@ def test_exists_hit_w_alternate_client(self): conn2 = _make_connection({}) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client1) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: self.assertTrue(job.exists(client=client2)) - final_attributes.assert_called_with({'path': "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, None, - job) + final_attributes.assert_called_with( + {"path": "/projects/{}/jobs/{}".format(self.PROJECT, self.JOB_ID)}, + None, + job, + ) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -2598,10 +2647,14 @@ def test_exists_miss_w_job_reference(self): conn = _make_connection() client = _make_client(project=self.PROJECT, connection=conn) load_job = self._make_one(job_ref, [self.SOURCE1], self.TABLE_REF, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: self.assertFalse(load_job.exists()) - final_attributes.assert_called_with({'path': "/projects/other-project/jobs/my-job-id"}, None, load_job) + final_attributes.assert_called_with( + {"path": "/projects/other-project/jobs/my-job-id"}, None, load_job + ) conn.api_request.assert_called_once_with( method="GET", @@ -2616,10 +2669,12 @@ def test_reload_w_bound_client(self): conn = _make_connection(RESOURCE) client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload() - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None @@ -2634,10 +2689,12 @@ def test_reload_w_alternate_client(self): conn2 = _make_connection(RESOURCE) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client1) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -2655,11 +2712,16 @@ def test_reload_w_job_reference(self): conn = _make_connection(resource) client = _make_client(project=self.PROJECT, connection=conn) load_job = self._make_one(job_ref, [self.SOURCE1], self.TABLE_REF, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: load_job.reload() - final_attributes.assert_called_with({'path': "/projects/alternative-project/jobs/{}".format(self.JOB_ID)}, None, - load_job) + final_attributes.assert_called_with( + {"path": "/projects/alternative-project/jobs/{}".format(self.JOB_ID)}, + None, + load_job, + ) conn.api_request.assert_called_once_with( method="GET", @@ -2675,10 +2737,12 @@ def test_cancel_w_bound_client(self): conn = _make_connection(RESPONSE) client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.cancel() - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", path=PATH, query_params={}, timeout=None, @@ -2694,10 +2758,12 @@ def test_cancel_w_alternate_client(self): conn2 = _make_connection(RESPONSE) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client1) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.cancel(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -2715,11 +2781,20 @@ def test_cancel_w_job_reference(self): conn = _make_connection({"job": resource}) client = _make_client(project=self.PROJECT, connection=conn) load_job = self._make_one(job_ref, [self.SOURCE1], self.TABLE_REF, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: load_job.cancel() final_attributes.assert_called_with( - {'path': "/projects/alternative-project/jobs/{}/cancel".format(self.JOB_ID)}, None, load_job) + { + "path": "/projects/alternative-project/jobs/{}/cancel".format( + self.JOB_ID + ) + }, + None, + load_job, + ) conn.api_request.assert_called_once_with( method="POST", path="/projects/alternative-project/jobs/{}/cancel".format(self.JOB_ID), @@ -3014,10 +3089,12 @@ def test_begin_w_bound_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -3080,10 +3157,12 @@ def test_begin_w_alternate_client(self): config.create_disposition = CreateDisposition.CREATE_NEVER config.write_disposition = WriteDisposition.WRITE_TRUNCATE job = self._make_one(self.JOB_ID, [source], destination, client1, config) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job._begin(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3105,10 +3184,12 @@ def test_exists_miss_w_bound_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None, @@ -3123,10 +3204,12 @@ def test_exists_hit_w_alternate_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client1) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) 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}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3141,10 +3224,12 @@ def test_reload_w_bound_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload() - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None @@ -3161,10 +3246,12 @@ def test_reload_w_alternate_client(self): source = self._table_ref(self.SOURCE_TABLE) destination = self._table_ref(self.DESTINATION_TABLE) job = self._make_one(self.JOB_ID, [source], destination, client1) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3424,10 +3511,12 @@ def test_begin_w_bound_client(self): source_dataset = DatasetReference(self.PROJECT, self.DS_ID) source = source_dataset.table(self.SOURCE_TABLE) job = self._make_one(self.JOB_ID, source, [self.DESTINATION_URI], client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -3484,10 +3573,12 @@ def test_begin_w_alternate_client(self): job = self._make_one( self.JOB_ID, source, [self.DESTINATION_URI], client1, config ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job._begin(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3508,10 +3599,12 @@ def test_exists_miss_w_bound_client(self): job = self._make_one( self.JOB_ID, self.TABLE_REF, [self.DESTINATION_URI], client ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None, @@ -3526,10 +3619,12 @@ def test_exists_hit_w_alternate_client(self): job = self._make_one( self.JOB_ID, self.TABLE_REF, [self.DESTINATION_URI], client1 ) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) 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}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -3546,10 +3641,12 @@ def test_reload_w_bound_client(self): source_dataset = DatasetReference(self.PROJECT, self.DS_ID) source = source_dataset.table(self.SOURCE_TABLE) job = self._make_one(self.JOB_ID, source, [self.DESTINATION_URI], client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload() - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={}, timeout=None ) @@ -3567,10 +3664,12 @@ def test_reload_w_alternate_client(self): source_dataset = DatasetReference(self.PROJECT, self.DS_ID) source = source_dataset.table(self.SOURCE_TABLE) job = self._make_one(self.JOB_ID, source, [self.DESTINATION_URI], client1) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -4109,7 +4208,7 @@ def test_done_w_timeout(self): job = self._get_target_class().from_api_repr(resource, client) with mock.patch.object( - client, "_get_query_results" + client, "_get_query_results" ) as fake_get_results, mock.patch.object(job, "reload") as fake_reload: job.done(timeout=42) @@ -4127,7 +4226,7 @@ def test_done_w_timeout_and_longer_internal_api_timeout(self): job._done_timeout = 8.8 with mock.patch.object( - client, "_get_query_results" + client, "_get_query_results" ) as fake_get_results, mock.patch.object(job, "reload") as fake_reload: job.done(timeout=5.5) @@ -4909,10 +5008,12 @@ def test__begin_w_timeout(self): conn = _make_connection(RESOURCE) client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, self.QUERY, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job._begin(timeout=7.5) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -4944,10 +5045,12 @@ def test_begin_w_bound_client(self): config = QueryJobConfig() config.default_dataset = DatasetReference(self.PROJECT, DS_ID) job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=config) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) self.assertIsNone(job.default_dataset) self.assertEqual(job.udf_resources, []) @@ -5026,10 +5129,12 @@ def test_begin_w_alternate_client(self): config.maximum_bytes_billed = 123456 config.schema_update_options = [SchemaUpdateOption.ALLOW_FIELD_RELAXATION] job = self._make_one(self.JOB_ID, self.QUERY, client1, job_config=config) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job._begin(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5070,10 +5175,12 @@ def test_begin_w_udf(self): config.udf_resources = udf_resources config.use_legacy_sql = True job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=config) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) self.assertEqual(job.udf_resources, udf_resources) conn.api_request.assert_called_once_with( @@ -5122,10 +5229,12 @@ def test_begin_w_named_query_parameter(self): jconfig = QueryJobConfig() jconfig.query_parameters = query_parameters job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=jconfig) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) self.assertEqual(job.query_parameters, query_parameters) conn.api_request.assert_called_once_with( @@ -5168,10 +5277,12 @@ def test_begin_w_positional_query_parameter(self): jconfig = QueryJobConfig() jconfig.query_parameters = query_parameters job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=jconfig) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) self.assertEqual(job.query_parameters, query_parameters) conn.api_request.assert_called_once_with( @@ -5246,10 +5357,12 @@ def test_begin_w_table_defs(self): config.table_definitions = {bt_table: bt_config, csv_table: csv_config} config.use_legacy_sql = True job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=config) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="POST", @@ -5287,10 +5400,12 @@ def test_dry_run_query(self): config = QueryJobConfig() config.dry_run = True job = self._make_one(self.JOB_ID, self.QUERY, client, job_config=config) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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) + final_attributes.assert_called_with({"path": PATH}, None, job) self.assertEqual(job.udf_resources, []) conn.api_request.assert_called_once_with( method="POST", @@ -5311,10 +5426,12 @@ def test_exists_miss_w_bound_client(self): conn = _make_connection() client = _make_client(project=self.PROJECT, connection=conn) job = self._make_one(self.JOB_ID, self.QUERY, client) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + 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': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn.api_request.assert_called_once_with( method="GET", path=PATH, query_params={"fields": "id"}, timeout=None @@ -5327,10 +5444,12 @@ def test_exists_hit_w_alternate_client(self): conn2 = _make_connection({}) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, self.QUERY, client1) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) 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}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5352,10 +5471,12 @@ def test_reload_w_bound_client(self): config = QueryJobConfig() config.destination = table_ref job = self._make_one(self.JOB_ID, None, client, job_config=config) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload() - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) self.assertNotEqual(job.destination, table_ref) @@ -5380,10 +5501,12 @@ def test_reload_w_alternate_client(self): conn2 = _make_connection(RESOURCE) client2 = _make_client(project=self.PROJECT, connection=conn2) job = self._make_one(self.JOB_ID, self.QUERY, client1) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload(client=client2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) conn1.api_request.assert_not_called() conn2.api_request.assert_called_once_with( @@ -5406,10 +5529,12 @@ def test_reload_w_timeout(self): config = QueryJobConfig() config.destination = table_ref job = self._make_one(self.JOB_ID, None, client, job_config=config) - with mock.patch('google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes') as final_attributes: + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: job.reload(timeout=4.2) - final_attributes.assert_called_with({'path': PATH}, None, job) + final_attributes.assert_called_with({"path": PATH}, None, job) self.assertNotEqual(job.destination, table_ref) @@ -6130,17 +6255,17 @@ def test_from_api_repr_normal(self): @pytest.mark.parametrize( "query,expected", ( - (None, False), - ("", False), - ("select name, age from table", False), - ("select name, age from table LIMIT 10;", False), - ("select name, age from table order by other_column;", True), - ("Select name, age From table Order By other_column", True), - ("SELECT name, age FROM table ORDER BY other_column;", True), - ("select name, age from table order\nby other_column", True), - ("Select name, age From table Order\nBy other_column;", True), - ("SELECT name, age FROM table ORDER\nBY other_column", True), - ("SelecT name, age froM table OrdeR \n\t BY other_column;", True), + (None, False), + ("", False), + ("select name, age from table", False), + ("select name, age from table LIMIT 10;", False), + ("select name, age from table order by other_column;", True), + ("Select name, age From table Order By other_column", True), + ("SELECT name, age FROM table ORDER BY other_column;", True), + ("select name, age from table order\nby other_column", True), + ("Select name, age From table Order\nBy other_column;", True), + ("SELECT name, age FROM table ORDER\nBY other_column", True), + ("SelecT name, age froM table OrdeR \n\t BY other_column;", True), ), ) def test__contains_order_by(query, expected): @@ -6159,13 +6284,13 @@ def test__contains_order_by(query, expected): @pytest.mark.parametrize( "query", ( - "select name, age from table order by other_column;", - "Select name, age From table Order By other_column;", - "SELECT name, age FROM table ORDER BY other_column;", - "select name, age from table order\nby other_column;", - "Select name, age From table Order\nBy other_column;", - "SELECT name, age FROM table ORDER\nBY other_column;", - "SelecT name, age froM table OrdeR \n\t BY other_column;", + "select name, age from table order by other_column;", + "Select name, age From table Order By other_column;", + "SELECT name, age FROM table ORDER BY other_column;", + "select name, age from table order\nby other_column;", + "Select name, age From table Order\nBy other_column;", + "SELECT name, age FROM table ORDER\nBY other_column;", + "SelecT name, age froM table OrdeR \n\t BY other_column;", ), ) def test_to_dataframe_bqstorage_preserve_order(query): From cd8e3033267076fe35f51e6bfd183b739cf908f0 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Mon, 17 Aug 2020 10:39:19 -0400 Subject: [PATCH 21/41] fixing trace not defined issue --- google/cloud/bigquery/opentelemetry_tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 44ed831bc..c87488cc2 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -39,7 +39,6 @@ "db.system": "BigQuery", } -tracer = trace.get_tracer(__name__) @contextmanager @@ -48,6 +47,7 @@ def create_span(name, attributes=None, client=None, job_ref=None): if not HAS_OPENTELEMETRY: yield None return + tracer = trace.get_tracer(__name__) final_attributes = _get_final_span_attributes(attributes, client, job_ref) with tracer.start_as_current_span(name=name, attributes=final_attributes) as span: From d8c601d3d705a807936258405143edb107e14743 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Mon, 17 Aug 2020 10:56:23 -0400 Subject: [PATCH 22/41] fixing lint issues --- google/cloud/bigquery/opentelemetry_tracing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index c87488cc2..ddfed0cec 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -40,7 +40,6 @@ } - @contextmanager def create_span(name, attributes=None, client=None, job_ref=None): # yield new span value From 2bcd106876db8b552b274b94d9497ea5c98f8cd7 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Mon, 17 Aug 2020 18:05:58 -0400 Subject: [PATCH 23/41] fixing documentation issues and python2 testing issue --- .../cloud/bigquery/opentelemetry_tracing.py | 36 ++++---- noxfile.py | 2 +- tests/unit/test_opentelemetry_tracing.py | 91 +++++++++---------- 3 files changed, 62 insertions(+), 67 deletions(-) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index ddfed0cec..063ec3534 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -34,7 +34,6 @@ HAS_OPENTELEMETRY = False - _default_attributes = { "db.system": "BigQuery", } @@ -42,24 +41,6 @@ @contextmanager def create_span(name, attributes=None, client=None, job_ref=None): - # yield new span value - if not HAS_OPENTELEMETRY: - yield None - return - tracer = trace.get_tracer(__name__) - - final_attributes = _get_final_span_attributes(attributes, client, job_ref) - with tracer.start_as_current_span(name=name, attributes=final_attributes) as span: - try: - yield span - span.set_status(Status(http_status_to_canonical_code(200))) - except GoogleAPICallError as error: - if error.code is not None: - span.set_status(Status(http_status_to_canonical_code(error.code))) - raise - - -def _get_final_span_attributes(attributes=None, client=None, job_ref=None): """Creates a ContextManager for a Span to be exported to the configured exporter. If no configuration exists yields None. @@ -83,7 +64,24 @@ def _get_final_span_attributes(attributes=None, client=None, job_ref=None): Raised if a span could not be yielded or issue with call to OpenTelemetry. """ + final_attributes = _get_final_span_attributes(attributes, client, job_ref) + if not HAS_OPENTELEMETRY: + yield None + return + tracer = trace.get_tracer(__name__) + + # yield new span value + with tracer.start_as_current_span(name=name, attributes=final_attributes) as span: + try: + yield span + span.set_status(Status(http_status_to_canonical_code(200))) + except GoogleAPICallError as error: + if error.code is not None: + span.set_status(Status(http_status_to_canonical_code(error.code))) + raise + +def _get_final_span_attributes(attributes=None, client=None, job_ref=None): if client: client_attributes = _set_client_attributes(client) _default_attributes.update(client_attributes) diff --git a/noxfile.py b/noxfile.py index c4c0ceb7b..2de3f23c1 100644 --- a/noxfile.py +++ b/noxfile.py @@ -33,7 +33,7 @@ def default(session): """ # Install all test dependencies, then install local packages in-place. session.install( - "mock", "pytest", "google-cloud-testutils", "pytest-cov", "freezegun", + "mock", "pytest", "google-cloud-testutils", "pytest-cov", "freezegun" ) session.install("grpcio") diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 18da91536..1e13f2dad 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -1,20 +1,26 @@ import sys import pytest -from importlib import reload import mock from google.cloud.bigquery import opentelemetry_tracing, client -from opentelemetry import trace -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +try: + import opentelemetry + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +except ImportError: + opentelemetry = None + +from six.moves import reload_module TEST_SPAN_NAME = "bar" TEST_SPAN_ATTRIBUTES = {"foo": "baz"} +@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") @pytest.fixture def setup(): tracer_provider = TracerProvider() @@ -25,35 +31,31 @@ def setup(): yield memory_exporter +@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_opentelemetry_not_installed(setup): temp_module = sys.modules["opentelemetry"] sys.modules["opentelemetry"] = None - reload(opentelemetry_tracing) + reload_module(opentelemetry_tracing) with opentelemetry_tracing.create_span("No-op for opentelemetry") as span: assert span is None sys.modules["opentelemetry"] = temp_module - reload(opentelemetry_tracing) + reload_module(opentelemetry_tracing) +@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_opentelemetry_success(setup): expected_attributes = {"foo": "baz", "db.system": "BigQuery"} + with opentelemetry_tracing.create_span( - TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=None, job_ref=None ) as span: - if span is None: - span_list = setup.get_finished_spans() - assert len(span_list) == 1 + assert span is not None assert span.name == TEST_SPAN_NAME assert span.attributes == expected_attributes +@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_default_client_attributes(setup): - import google.auth.credentials - - mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) - test_client = client.Client( - project="test_project", credentials=mock_credentials, location="test_location" - ) expected_attributes = { "foo": "baz", @@ -61,29 +63,19 @@ def test_default_client_attributes(setup): "db.name": "test_project", "location": "test_location", } - with opentelemetry_tracing.create_span( - TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client - ) as span: - if span is None: - span_list = setup.get_finished_spans() - assert len(span_list) == 1 - assert span.name == TEST_SPAN_NAME - assert span.attributes == expected_attributes - - + with mock.patch('google.cloud.bigquery.client.Client') as test_client: + test_client.project = 'test_project' + test_client.location = 'test_location' + with opentelemetry_tracing.create_span( + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client + ) as span: + assert span is not None + assert span.name == TEST_SPAN_NAME + assert span.attributes == expected_attributes + + +@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_default_job_attributes(setup): - from google.cloud.bigquery import job - import google.auth.credentials - - mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) - - test_job_reference = job._JobReference( - job_id="test_job_id", project="test_project_id", location="test_location" - ) - test_client = client.Client( - project="test_project", credentials=mock_credentials, location="test_location" - ) - test_job = job._AsyncJob(job_id=test_job_reference, client=test_client) expected_attributes = { "db.system": "BigQuery", @@ -93,12 +85,17 @@ def test_default_job_attributes(setup): "job_id": "test_job_id", "foo": "baz", } + with mock.patch('google.cloud.bigquery.job._AsyncJob') as test_job_ref: + test_job_ref.job_id = 'test_job_id' + test_job_ref.location = 'test_location' + test_job_ref.project = 'test_project_id' + test_job_ref.num_child_jobs = '0' + + with opentelemetry_tracing.create_span( + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job_ref + ) as span: + assert span is not None + assert span.name == TEST_SPAN_NAME + assert span.attributes == expected_attributes + - with opentelemetry_tracing.create_span( - TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job - ) as span: - if span is None: - span_list = setup.get_finished_spans() - assert len(span_list) == 1 - assert span.name == TEST_SPAN_NAME - assert span.attributes == expected_attributes From f1de589c99f0bd6c2f3c3335b5c486df4fabe5b1 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Mon, 17 Aug 2020 22:46:04 -0400 Subject: [PATCH 24/41] linting and fixing coverage issues --- tests/unit/test_client.py | 38 ++++++++++++++++-------- tests/unit/test_opentelemetry_tracing.py | 30 +++++++++---------- 2 files changed, 40 insertions(+), 28 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 4166e76df..14ce8e2de 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -258,7 +258,9 @@ def test__get_query_results_miss_w_explicit_project_and_timeout(self): timeout_ms=500, timeout=42, ) - final_attributes.assert_called_once_with({"path": path}, client, None) + final_attributes.assert_called_once_with( + {"path": path}, client, None + ) # pragma: NO COVER conn.api_request.assert_called_once_with( method="GET", @@ -673,7 +675,7 @@ def test_get_dataset(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with( + final_attributes.assert_called_once_with( # pragma: NO COVER {"path": "/%s" % path}, client, None ) @@ -684,7 +686,7 @@ def test_get_dataset(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with( + final_attributes.assert_called_once_with( # pragma: NO COVER {"path": "/%s" % path}, client, None ) @@ -698,7 +700,9 @@ def test_get_dataset(self): ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with({"path": path}, client, None) + final_attributes.assert_called_once_with( + {"path": path}, client, None + ) # pragma: NO COVER # Retryable reason, but retry is disabled. client._connection = make_connection( @@ -710,7 +714,7 @@ def test_get_dataset(self): ) as final_attributes: client.get_dataset(dataset_ref, retry=None) - final_attributes.assert_called_once_with( + final_attributes.assert_called_once_with( # pragma: NO COVER {"path": "/%s" % path}, client, None ) @@ -1157,7 +1161,9 @@ def test_create_dataset_alreadyexists_w_exists_ok_false(self): ) as final_attributes: client.create_dataset(self.DS_ID) - final_attributes.assert_called_once_with({"path": "path"}, client, None) + final_attributes.assert_called_once_with( + {"path": "path"}, client, None + ) # pragma: NO COVER def test_create_dataset_alreadyexists_w_exists_ok_true(self): post_path = "/projects/{}/datasets".format(self.PROJECT) @@ -1257,7 +1263,7 @@ def test_create_routine_w_conflict(self): ) as final_attributes: client.create_routine(routine) - final_attributes.assert_called_once_with( + final_attributes.assert_called_once_with( # pragma: NO COVER {"path": path, "exists_okay": False}, client, None ) @@ -1721,7 +1727,7 @@ def test_create_table_alreadyexists_w_exists_ok_false(self): ) as final_attributes: client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_with( + final_attributes.assert_called_with( # pragma: NO COVER {"path": post_path, "dataset_id": self.TABLE_REF.dataset_id}, client, None, @@ -2137,7 +2143,7 @@ def test_set_iam_policy_invalid_policy(self): ) as final_attributes: client.set_iam_policy(self.TABLE_REF, invalid_policy_repr) - final_attributes.assert_called_once_with( + final_attributes.assert_called_once_with( # pragma: NO COVER {"path": "{}:setIamPolicy".format(self.TABLE_REF.path)}, client, None ) @@ -2162,7 +2168,7 @@ def test_set_iam_policy_w_invalid_table(self): ) as final_attributes: client.set_iam_policy(table_resource_string, policy) - final_attributes.assert_called_once_with( + final_attributes.assert_called_once_with( # pragma: NO COVER {"path": table_resource_string}, client, None ) @@ -3246,7 +3252,9 @@ def test_delete_dataset_w_not_found_ok_false(self): ) as final_attributes: client.delete_dataset(self.DS_ID) - final_attributes.assert_called_once_with({"path": path}, client, None) + final_attributes.assert_called_once_with( + {"path": path}, client, None + ) # pragma: NO COVER conn.api_request.assert_called_with( method="DELETE", path=path, query_params={}, timeout=None @@ -3397,7 +3405,9 @@ def test_delete_routine_w_not_found_ok_false(self): ) as final_attributes: client.delete_routine("routines-project.test_routines.test_routine") - final_attributes.assert_called_once_with({"path": path}, client, None) + final_attributes.assert_called_once_with( + {"path": path}, client, None + ) # pragma: NO COVER conn.api_request.assert_called_with( method="DELETE", path=path, timeout=None, @@ -3483,7 +3493,9 @@ def test_delete_table_w_not_found_ok_false(self): ) as final_attributes: client.delete_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_once_with({"path": path}, client, None) + final_attributes.assert_called_once_with( + {"path": path}, client, None + ) # pragma: NO COVER conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 1e13f2dad..a82674c29 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -3,14 +3,16 @@ import mock -from google.cloud.bigquery import opentelemetry_tracing, client +from google.cloud.bigquery import opentelemetry_tracing try: import opentelemetry from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor - from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) except ImportError: opentelemetry = None @@ -47,7 +49,7 @@ def test_opentelemetry_success(setup): expected_attributes = {"foo": "baz", "db.system": "BigQuery"} with opentelemetry_tracing.create_span( - TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=None, job_ref=None + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=None, job_ref=None ) as span: assert span is not None assert span.name == TEST_SPAN_NAME @@ -63,11 +65,11 @@ def test_default_client_attributes(setup): "db.name": "test_project", "location": "test_location", } - with mock.patch('google.cloud.bigquery.client.Client') as test_client: - test_client.project = 'test_project' - test_client.location = 'test_location' + with mock.patch("google.cloud.bigquery.client.Client") as test_client: + test_client.project = "test_project" + test_client.location = "test_location" with opentelemetry_tracing.create_span( - TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client ) as span: assert span is not None assert span.name == TEST_SPAN_NAME @@ -85,17 +87,15 @@ def test_default_job_attributes(setup): "job_id": "test_job_id", "foo": "baz", } - with mock.patch('google.cloud.bigquery.job._AsyncJob') as test_job_ref: - test_job_ref.job_id = 'test_job_id' - test_job_ref.location = 'test_location' - test_job_ref.project = 'test_project_id' - test_job_ref.num_child_jobs = '0' + with mock.patch("google.cloud.bigquery.job._AsyncJob") as test_job_ref: + test_job_ref.job_id = "test_job_id" + test_job_ref.location = "test_location" + test_job_ref.project = "test_project_id" + test_job_ref.num_child_jobs = "0" with opentelemetry_tracing.create_span( - TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job_ref + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job_ref ) as span: assert span is not None assert span.name == TEST_SPAN_NAME assert span.attributes == expected_attributes - - From 0c7b35492b6ada5f4cd5f082d5e1c02ebdf548bc Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Tue, 18 Aug 2020 19:20:38 -0400 Subject: [PATCH 25/41] adding suggested changes --- google/cloud/bigquery/client.py | 118 ++++--------- .../cloud/bigquery/opentelemetry_tracing.py | 14 +- noxfile.py | 4 +- tests/unit/test_client.py | 156 +++++------------- tests/unit/test_opentelemetry_tracing.py | 15 ++ 5 files changed, 89 insertions(+), 218 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 54711f0ea..df98e43ba 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -288,16 +288,11 @@ def list_projects( Iterator of :class:`~google.cloud.bigquery.client.Project` accessible to the current client. """ - path = "/projects" - span_attributes = {"path": path} - with create_span( - name="BigQuery.listProjects", attributes=span_attributes, client=self - ): - api_request = functools.partial(self._call_api, retry, timeout=timeout) + return page_iterator.HTTPIterator( client=self, - api_request=api_request, - path=path, + api_request=functools.partial(self._call_api, retry, timeout=timeout), + path="/projects", item_to_value=_item_to_project, items_key="projects", page_token=page_token, @@ -357,17 +352,11 @@ def list_datasets( # TODO: consider supporting a dict of label -> value for filter, # and converting it into a string here. extra_params["filter"] = filter - path = "/projects/%s/datasets" % (project,) - span_attributes = {"path": path, "page_token": page_token} - with create_span( - name="BigQuery.listDatasets", attributes=span_attributes, client=self - ): - api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, - api_request=api_request, - path=path, + api_request=functools.partial(self._call_api, retry, timeout=timeout), + path="/projects/%s/datasets" % (project,), item_to_value=_item_to_dataset, items_key="datasets", page_token=page_token, @@ -538,7 +527,7 @@ def create_routine( ) resource = routine.to_api_repr() try: - span_attributes = {"path": path, "exists_ok": exists_ok} + span_attributes = {"path": path} with create_span( name="BigQuery.createRoutine", attributes=span_attributes, client=self ): @@ -657,9 +646,9 @@ def get_iam_policy( body = {"options": {"requestedPolicyVersion": 1}} path = "{}:getIamPolicy".format(table.path) - span_attributes = {"path": path, "body": body} + span_attributes = {"path": path} with create_span( - name="BigQuery.getIAMPolicy", attributes=span_attributes, client=self + name="BigQuery.getIamPolicy", attributes=span_attributes, client=self ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, @@ -684,7 +673,7 @@ def set_iam_policy( path = "{}:setIamPolicy".format(table.path) span_attributes = {"path": path} with create_span( - name="BigQuery.setIAMPolicy", attributes=span_attributes, client=self + name="BigQuery.setIamPolicy", attributes=span_attributes, client=self ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, @@ -701,9 +690,9 @@ def test_iam_permissions( body = {"permissions": permissions} path = "{}:testIamPermissions".format(table.path) - span_attributes = {"path": path, "body": body} + span_attributes = {"path": path} with create_span( - name="BigQuery.testIAMPermissions", attributes=span_attributes, client=self + name="BigQuery.testIamPermissions", attributes=span_attributes, client=self ): response = self._call_api( retry, method="POST", path=path, data=body, timeout=timeout, @@ -853,7 +842,7 @@ def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): else: headers = None path = dataset.path - span_attributes = {"path": path, "fields": fields} + span_attributes = {"path": path, "fields": json.dumps(fields)} with create_span( name="BigQuery.updateDataset", attributes=span_attributes, client=self @@ -902,7 +891,7 @@ def update_model(self, model, fields, retry=DEFAULT_RETRY, timeout=None): else: headers = None path = model.path - span_attributes = {"path": path, "fields": fields} + span_attributes = {"path": path, "fields": json.dumps(fields)} with create_span( name="BigQuery.updateModel", attributes=span_attributes, client=self @@ -962,7 +951,7 @@ def update_routine(self, routine, fields, retry=DEFAULT_RETRY, timeout=None): partial["routineReference"] = routine.reference.to_api_repr() path = routine.path - span_attributes = {"path": path, "fields": fields} + span_attributes = {"path": path, "fields": json.dumps(fields)} with create_span( name="BigQuery.updateRoutine", attributes=span_attributes, client=self @@ -1012,7 +1001,7 @@ def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): headers = None path = table.path - span_attributes = {"path": path, "fields": fields} + span_attributes = {"path": path, "fields": json.dumps(fields)} with create_span( name="BigQuery.updateTable", attributes=span_attributes, client=self @@ -1079,20 +1068,10 @@ def list_models( if not isinstance(dataset, (Dataset, DatasetReference)): raise TypeError("dataset must be a Dataset, DatasetReference, or string") - path = "%s/models" % dataset.path - span_attributes = { - "path": path, - "page_token": page_token, - "max_results": max_results, - } - with create_span( - name="BigQuery.listModels", attributes=span_attributes, client=self - ): - api_request = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( client=self, - api_request=api_request, - path=path, + api_request=functools.partial(self._call_api, retry, timeout=timeout), + path="%s/models" % dataset.path, item_to_value=_item_to_model, items_key="models", page_token=page_token, @@ -1153,20 +1132,10 @@ def list_routines( if not isinstance(dataset, (Dataset, DatasetReference)): raise TypeError("dataset must be a Dataset, DatasetReference, or string") - path = "{}/routines".format(dataset.path) - span_attributes = { - "path": path, - "page_token": page_token, - "max_results": max_results, - } - with create_span( - name="BigQuery.listRoutines", attributes=span_attributes, client=self - ): - api_requests = functools.partial(self._call_api, retry, timeout=timeout) result = page_iterator.HTTPIterator( client=self, - api_request=api_requests, - path=path, + api_request=functools.partial(self._call_api, retry, timeout=timeout), + path="{}/routines".format(dataset.path), item_to_value=_item_to_routine, items_key="routines", page_token=page_token, @@ -1227,21 +1196,10 @@ def list_tables( if not isinstance(dataset, (Dataset, DatasetReference)): raise TypeError("dataset must be a Dataset, DatasetReference, or string") - path = "%s/tables" % dataset.path - span_attributes = { - "path": path, - "page_token": page_token, - "max_results": max_results, - } - with create_span( - name="BigQuery.listTables", attributes=span_attributes, client=self - ): - api_requests = functools.partial(self._call_api, retry, timeout=timeout) - result = page_iterator.HTTPIterator( client=self, - api_request=api_requests, - path=path, + api_request=functools.partial(self._call_api, retry, timeout=timeout), + path="%s/tables" % dataset.path, item_to_value=_item_to_table, items_key="tables", page_token=page_token, @@ -1298,9 +1256,11 @@ def delete_dataset( path = dataset.path if delete_contents: params["deleteContents"] = "true" + span_attributes = {"path": path, "deleteContents": delete_contents} + else: + span_attributes = {"path": path} try: - span_attributes = {"path": path} with create_span( name="BigQuery.deleteDataset", attributes=span_attributes, client=self ): @@ -1803,21 +1763,10 @@ def list_jobs( if project is None: project = self.project - path = "/projects/%s/jobs" % (project,) - span_attributes = { - "path": path, - "all_users": all_users, - "state_filter": state_filter, - } - - with create_span( - name="BigQuery.listJobs", attributes=span_attributes, client=self - ): - api_request = functools.partial(self._call_api, retry, timeout=timeout) return page_iterator.HTTPIterator( client=self, - api_request=api_request, - path=path, + api_request=functools.partial(self._call_api, retry, timeout=timeout), + path="/projects/%s/jobs" % (project,), item_to_value=_item_to_job, items_key="jobs", page_token=page_token, @@ -3045,22 +2994,11 @@ def list_rows( params["selectedFields"] = ",".join(field.name for field in selected_fields) if start_index is not None: params["startIndex"] = start_index - path = "%s/data" % (table.path,) - span_attributes = { - "path": path, - "page_size": page_size, - "page_token": page_token, - "start_index": start_index, - } - with create_span( - name="BigQuery.listRows", attributes=span_attributes, client=self - ): - api_request = functools.partial(self._call_api, retry, timeout=timeout) row_iterator = RowIterator( client=self, - api_request=api_request, - path=path, + api_request=functools.partial(self._call_api, retry, timeout=timeout), + path="%s/data" % (table.path,), schema=schema, page_token=page_token, max_results=max_results, diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 063ec3534..90ab99666 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -1,4 +1,4 @@ -# Copyright 2020 Google LLC All rights reserved. +# Copyright 2020 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,10 +26,10 @@ except ImportError: Logger.info( - "This service is instrumented using opentelemetry." - "Opentelemetry could not be imported, please" + "This service is instrumented using OpenTelemetry." + "Opentelemetry could not be imported; please" "add opentelemetry-api and opentelemetry-instrumentation" - "packages, in order to get Big Query Tracing data" + "packages in order to get Big Query Tracing data." ) HAS_OPENTELEMETRY = False @@ -46,13 +46,13 @@ def create_span(name, attributes=None, client=None, job_ref=None): Args: name (str): Name that will be set for the span being created - attributes(Optional[dict]): + attributes (Optional[dict]): Additional attributes that pertain to the specific API call (i.e. not a default attribute) client (Optional[google.cloud.bigquery.client.Client]): Pass in a Client object to extract any attributes that may be relevant to it and add them to the created spans. - job_ref(Optional[google.cloud.bigquery.job._AsyncJob]) + job_ref (Optional[google.cloud.bigquery.job._AsyncJob]) Pass in a _AsyncJob object to extract any attributes that may be relevant to it and add them to the created spans. @@ -85,7 +85,7 @@ def _get_final_span_attributes(attributes=None, client=None, job_ref=None): if client: client_attributes = _set_client_attributes(client) _default_attributes.update(client_attributes) - elif job_ref: + if job_ref: job_attributes = _set_job_attributes(job_ref) _default_attributes.update(job_attributes) diff --git a/noxfile.py b/noxfile.py index 2de3f23c1..e4aad1866 100644 --- a/noxfile.py +++ b/noxfile.py @@ -48,8 +48,8 @@ def default(session): else: session.install("ipython") - # opentelemetry was not added to [all] because opentelemetry does not support Python 2 - # exporter does not need to be in nox thus it has been added to README documentation + # opentelemetry was not added to [all] because opentelemetry does not support Python 2. + # Exporter does not need to be in nox thus it has been added to README documentation if session.python != "2.7": session.install( "opentelemetry-api==0.9b0", diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 14ce8e2de..3d0caef04 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -412,11 +412,7 @@ def test_list_projects_defaults(self): creds = _make_credentials() client = self._make_one(PROJECT_1, creds) conn = client._connection = make_connection(DATA) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_projects() - final_attributes.assert_called_once_with({"path": "/projects"}, client, None) + iterator = client.list_projects() page = six.next(iterator.pages) projects = list(page) token = iterator.next_page_token @@ -443,11 +439,8 @@ def test_list_projects_w_timeout(self): creds = _make_credentials() client = self._make_one(PROJECT_1, creds) conn = client._connection = make_connection(DATA) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_projects(timeout=7.5) - final_attributes.assert_called_once_with({"path": "/projects"}, client, None) + + iterator = client.list_projects(timeout=7.5) six.next(iterator.pages) @@ -462,12 +455,7 @@ def test_list_projects_explicit_response_missing_projects_key(self): client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_projects(max_results=3, page_token=TOKEN) - final_attributes.assert_called_once_with({"path": "/projects"}, client, None) - + iterator = client.list_projects(max_results=3, page_token=TOKEN) page = six.next(iterator.pages) projects = list(page) token = iterator.next_page_token @@ -515,13 +503,7 @@ def test_list_datasets_defaults(self): creds = _make_credentials() client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_datasets() - final_attributes.assert_called_once_with( - {"path": "/%s" % PATH, "page_token": None}, client, None - ) + iterator = client.list_datasets() page = six.next(iterator.pages) datasets = list(page) @@ -543,12 +525,7 @@ def test_list_datasets_w_project_and_timeout(self): client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection({}) - 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() + list(client.list_datasets(project="other-project", timeout=7.5)) conn.api_request.assert_called_once_with( method="GET", @@ -566,14 +543,8 @@ def test_list_datasets_explicit_response_missing_datasets_key(self): client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_datasets( - include_all=True, filter=FILTER, max_results=3, page_token=TOKEN - ) - final_attributes.assert_called_once_with( - {"path": "/%s" % PATH, "page_token": TOKEN}, client, None + iterator = client.list_datasets( + include_all=True, filter=FILTER, max_results=3, page_token=TOKEN ) page = six.next(iterator.pages) @@ -1235,7 +1206,7 @@ def test_create_routine_w_minimal_resource(self): actual_routine = client.create_routine(routine, timeout=7.5) final_attributes.assert_called_once_with( - {"path": path, "exists_ok": False}, client, None + {"path": path}, client, None ) conn.api_request.assert_called_once_with( @@ -1658,6 +1629,7 @@ def test_create_table_w_fully_qualified_string(self): got = client.create_table( "{}.{}.{}".format(self.PROJECT, self.DS_ID, self.TABLE_ID) ) + final_attributes.assert_called_once_with( {"path": "/%s" % path, "dataset_id": self.TABLE_REF.dataset_id}, client, @@ -2014,7 +1986,7 @@ def test_get_iam_policy(self): policy = client.get_iam_policy(self.TABLE_REF, timeout=7.5) final_attributes.assert_called_once_with( - {"path": PATH, "body": BODY}, client, None + {"path": PATH}, client, None ) conn.api_request.assert_called_once_with( @@ -2193,7 +2165,7 @@ def test_test_iam_permissions(self): client.test_iam_permissions(self.TABLE_REF, PERMISSIONS, timeout=7.5) final_attributes.assert_called_once_with( - {"path": PATH, "body": BODY}, client, None + {"path": PATH}, client, None ) conn.api_request.assert_called_once_with( @@ -2269,7 +2241,7 @@ def test_update_dataset(self): ) as final_attributes: ds2 = client.update_dataset(ds, fields=fields, timeout=7.5,) final_attributes.assert_called_once_with( - {"path": "/" + PATH, "fields": fields}, client, None + {"path": "/" + PATH, "fields": json.dumps(fields)}, client, None ) conn.api_request.assert_called_once_with( @@ -2319,7 +2291,7 @@ def test_update_dataset_w_custom_property(self): dataset = client.update_dataset(dataset, ["newAlphaProperty"]) final_attributes.assert_called_once_with( - {"path": path, "fields": ["newAlphaProperty"]}, client, None + {"path": path, "fields": json.dumps(["newAlphaProperty"])}, client, None ) conn.api_request.assert_called_once_with( @@ -2374,7 +2346,7 @@ def test_update_model(self): ) as final_attributes: updated_model = client.update_model(model, fields, timeout=7.5) final_attributes.assert_called_once_with( - {"path": "/" + path, "fields": fields}, client, None + {"path": "/" + path, "fields": json.dumps(fields)}, client, None ) sent = { @@ -2446,7 +2418,7 @@ def test_update_routine(self): ) as final_attributes: actual_routine = client.update_routine(routine, fields, timeout=7.5,) final_attributes.assert_called_once_with( - {"path": routine.path, "fields": fields}, client, None + {"path": routine.path, "fields": json.dumps(fields)}, client, None ) # TODO: routineReference isn't needed when the Routines API supports @@ -2472,7 +2444,7 @@ def test_update_routine(self): client.update_routine(routine, []) final_attributes.assert_called_once_with( - {"path": routine.path, "fields": []}, client, None + {"path": routine.path, "fields": json.dumps([])}, client, None ) req = conn.api_request.call_args @@ -2532,7 +2504,7 @@ def test_update_table(self): updated_table = client.update_table(table, fields, timeout=7.5) span_path = "/%s" % path final_attributes.assert_called_once_with( - {"path": span_path, "fields": fields}, client, None + {"path": span_path, "fields": json.dumps(fields)}, client, None ) sent = { @@ -2571,7 +2543,7 @@ def test_update_table(self): ) as final_attributes: client.update_table(table, []) final_attributes.assert_called_once_with( - {"path": "/" + path, "fields": []}, client, None + {"path": "/" + path, "fields": json.dumps([])}, client, None ) req = conn.api_request.call_args @@ -2599,7 +2571,7 @@ def test_update_table_w_custom_property(self): updated_table = client.update_table(table, ["newAlphaProperty"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": ["newAlphaProperty"]}, client, None + {"path": "/%s" % path, "fields": json.dumps(["newAlphaProperty"])}, client, None ) conn.api_request.assert_called_once_with( @@ -2634,7 +2606,7 @@ def test_update_table_only_use_legacy_sql(self): updated_table = client.update_table(table, ["view_use_legacy_sql"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": ["view_use_legacy_sql"]}, client, None + {"path": "/%s" % path, "fields": json.dumps(["view_use_legacy_sql"])}, client, None ) conn.api_request.assert_called_once_with( @@ -2704,7 +2676,7 @@ def test_update_table_w_query(self): updated_table = client.update_table(table, updated_properties) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": updated_properties}, client, None + {"path": "/%s" % path, "fields": json.dumps(updated_properties)}, client, None ) self.assertEqual(updated_table.schema, table.schema) @@ -2769,7 +2741,7 @@ def test_update_table_w_schema_None(self): updated_table = client.update_table(table, ["schema"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": ["schema"]}, client, None + {"path": "/%s" % path, "fields": json.dumps(["schema"])}, client, None ) self.assertEqual(len(conn.api_request.call_args_list), 2) @@ -2807,7 +2779,7 @@ def test_update_table_delete_property(self): table2 = client.update_table(table, ["description", "friendly_name"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": ["description", "friendly_name"]}, + {"path": "/%s" % path, "fields": json.dumps(["description", "friendly_name"])}, client, None, ) @@ -2821,7 +2793,7 @@ def test_update_table_delete_property(self): table3 = client.update_table(table2, ["description"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": ["description"]}, client, None + {"path": "/%s" % path, "fields": json.dumps(["description"])}, client, None ) self.assertEqual(len(conn.api_request.call_args_list), 2) @@ -2837,16 +2809,9 @@ def test_list_tables_empty_w_timeout(self): creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection({}) - dataset = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_tables(dataset, timeout=7.5) - final_attributes.assert_called_once_with( - {"path": path, "page_token": None, "max_results": None}, client, None - ) + iterator = client.list_tables(dataset, timeout=7.5) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2866,16 +2831,8 @@ def test_list_models_empty_w_timeout(self): conn = client._connection = make_connection({}) dataset_id = "{}.{}".format(self.PROJECT, self.DS_ID) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_models(dataset_id, timeout=7.5) - final_attributes.assert_called_once_with( - {"path": path, "page_token": iterator.next_page_token, "max_results": None}, - client, - None, - ) + iterator = client.list_models(dataset_id, timeout=7.5) page = six.next(iterator.pages) models = list(page) @@ -2919,16 +2876,7 @@ def test_list_models_defaults(self): conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_models(dataset) - - final_attributes.assert_called_once_with( - {"path": "/%s" % PATH, "page_token": None, "max_results": None}, - client, - None, - ) + iterator = client.list_models(dataset) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -2956,15 +2904,7 @@ def test_list_routines_empty_w_timeout(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection({}) path = "/projects/test-routines/datasets/test_routines/routines" - - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_routines("test-routines.test_routines", timeout=7.5) - - final_attributes.assert_called_once_with( - {"path": path, "page_token": None, "max_results": None}, client, None - ) + iterator = client.list_routines("test-routines.test_routines", timeout=7.5) page = six.next(iterator.pages) routines = list(page) @@ -3009,14 +2949,7 @@ def test_list_routines_defaults(self): client = self._make_one(project=project_id, credentials=creds) conn = client._connection = make_connection(resource) dataset = DatasetReference(client.project, dataset_id) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_routines(dataset) - - final_attributes.assert_called_once_with( - {"path": path, "page_token": None, "max_results": None}, client, None - ) + iterator = client.list_routines(dataset) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -3080,16 +3013,7 @@ def test_list_tables_defaults(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_tables(dataset) - - final_attributes.assert_called_once_with( - {"path": "/%s" % PATH, "page_token": None, "max_results": None}, - client, - None, - ) + iterator = client.list_tables(dataset) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) @@ -3143,17 +3067,11 @@ def test_list_tables_explicit(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - iterator = client.list_tables( - # Test with string for dataset ID. - self.DS_ID, - max_results=3, - page_token=TOKEN, - ) - final_attributes.assert_called_once_with( - {"path": "/%s" % PATH, "page_token": TOKEN, "max_results": 3}, client, None + iterator = client.list_tables( + # Test with string for dataset ID. + self.DS_ID, + max_results=3, + page_token=TOKEN, ) self.assertEqual(iterator.dataset, dataset) @@ -3220,7 +3138,7 @@ def test_delete_dataset_delete_contents(self): client.delete_dataset(arg, delete_contents=True) final_attributes.assert_called_once_with( - {"path": "/%s" % PATH}, client, None + {"path": "/%s" % PATH, "deleteContents": True}, client, None ) conn.api_request.assert_called_with( method="DELETE", diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index a82674c29..0064c9160 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -1,3 +1,18 @@ +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + import sys import pytest From ec96f5cba4040c9bcd05242a5eaa53bdc5ae195c Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Tue, 18 Aug 2020 19:44:26 -0400 Subject: [PATCH 26/41] linting --- tests/unit/test_client.py | 50 +++++++++++++++++++++++++++++---------- 1 file changed, 37 insertions(+), 13 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 3d0caef04..dd82f2075 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -258,6 +258,7 @@ def test__get_query_results_miss_w_explicit_project_and_timeout(self): timeout_ms=500, timeout=42, ) + final_attributes.assert_called_once_with( {"path": path}, client, None ) # pragma: NO COVER @@ -324,6 +325,7 @@ def test_get_service_account_email(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: service_account_email = client.get_service_account_email(timeout=7.5) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_once_with(method="GET", path=path, timeout=7.5) self.assertEqual(service_account_email, email) @@ -341,6 +343,7 @@ def test_get_service_account_email_w_alternate_project(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: service_account_email = client.get_service_account_email(project=project) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_once_with(method="GET", path=path, timeout=None) self.assertEqual(service_account_email, email) @@ -372,6 +375,7 @@ def test_get_service_account_email_w_custom_retry(self): service_account_email = client.get_service_account_email( retry=retry, timeout=7.5 ) + final_attributes.assert_called_once_with({"path": api_path}, client, None) self.assertEqual( service_account_email, "bq-123@bigquery-encryption.iam.gserviceaccount.com" @@ -646,6 +650,7 @@ def test_get_dataset(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: client.get_dataset(dataset_ref) + final_attributes.assert_called_once_with( # pragma: NO COVER {"path": "/%s" % path}, client, None ) @@ -657,6 +662,7 @@ def test_get_dataset(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: client.get_dataset(dataset_ref) + final_attributes.assert_called_once_with( # pragma: NO COVER {"path": "/%s" % path}, client, None ) @@ -700,6 +706,7 @@ def test_get_dataset(self): # Test with a string for dataset ID. dataset_ref.dataset_id ) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) @@ -831,6 +838,7 @@ def test_create_dataset_w_attrs(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: after = client.create_dataset(before) + final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) @@ -975,6 +983,7 @@ def test_create_dataset_w_client_location_w_dataset_location(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: after = client.create_dataset(before) + final_attributes.assert_called_once_with({"path": "/%s" % PATH}, client, None) self.assertEqual(after.dataset_id, self.DS_ID) @@ -1156,6 +1165,7 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: dataset = client.create_dataset(self.DS_ID, exists_ok=True) + final_attributes.assert_called_with({"path": get_path}, client, None) self.assertEqual(dataset.dataset_id, self.DS_ID) @@ -1205,9 +1215,7 @@ def test_create_routine_w_minimal_resource(self): ) as final_attributes: actual_routine = client.create_routine(routine, timeout=7.5) - final_attributes.assert_called_once_with( - {"path": path}, client, None - ) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_once_with( method="POST", path=path, data=resource, timeout=7.5, @@ -1345,6 +1353,7 @@ def test_create_table_w_custom_property(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: got = client.create_table(table) + final_attributes.assert_called_once_with( {"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None ) @@ -1739,6 +1748,7 @@ def test_create_table_alreadyexists_w_exists_ok_true(self): got = client.create_table( "{}.{}".format(self.DS_ID, self.TABLE_ID), exists_ok=True ) + final_attributes.assert_called_with({"path": get_path}, client, None) self.assertEqual(got.project, self.PROJECT) @@ -1985,9 +1995,7 @@ def test_get_iam_policy(self): ) as final_attributes: policy = client.get_iam_policy(self.TABLE_REF, timeout=7.5) - final_attributes.assert_called_once_with( - {"path": PATH}, client, None - ) + final_attributes.assert_called_once_with({"path": PATH}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -2062,6 +2070,7 @@ def test_set_iam_policy(self): returned_policy = client.set_iam_policy( self.TABLE_REF, policy, updateMask=MASK, timeout=7.5 ) + final_attributes.assert_called_once_with({"path": PATH}, client, None) conn.api_request.assert_called_once_with( @@ -2164,9 +2173,7 @@ def test_test_iam_permissions(self): ) as final_attributes: client.test_iam_permissions(self.TABLE_REF, PERMISSIONS, timeout=7.5) - final_attributes.assert_called_once_with( - {"path": PATH}, client, None - ) + final_attributes.assert_called_once_with({"path": PATH}, client, None) conn.api_request.assert_called_once_with( method="POST", path=PATH, data=BODY, timeout=7.5 @@ -2240,6 +2247,7 @@ def test_update_dataset(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: ds2 = client.update_dataset(ds, fields=fields, timeout=7.5,) + final_attributes.assert_called_once_with( {"path": "/" + PATH, "fields": json.dumps(fields)}, client, None ) @@ -2345,6 +2353,7 @@ def test_update_model(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: updated_model = client.update_model(model, fields, timeout=7.5) + final_attributes.assert_called_once_with( {"path": "/" + path, "fields": json.dumps(fields)}, client, None ) @@ -2417,6 +2426,7 @@ def test_update_routine(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: actual_routine = client.update_routine(routine, fields, timeout=7.5,) + final_attributes.assert_called_once_with( {"path": routine.path, "fields": json.dumps(fields)}, client, None ) @@ -2503,6 +2513,7 @@ def test_update_table(self): ) as final_attributes: updated_table = client.update_table(table, fields, timeout=7.5) span_path = "/%s" % path + final_attributes.assert_called_once_with( {"path": span_path, "fields": json.dumps(fields)}, client, None ) @@ -2542,6 +2553,7 @@ def test_update_table(self): "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: client.update_table(table, []) + final_attributes.assert_called_once_with( {"path": "/" + path, "fields": json.dumps([])}, client, None ) @@ -2571,7 +2583,9 @@ def test_update_table_w_custom_property(self): updated_table = client.update_table(table, ["newAlphaProperty"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(["newAlphaProperty"])}, client, None + {"path": "/%s" % path, "fields": json.dumps(["newAlphaProperty"])}, + client, + None, ) conn.api_request.assert_called_once_with( @@ -2606,7 +2620,9 @@ def test_update_table_only_use_legacy_sql(self): updated_table = client.update_table(table, ["view_use_legacy_sql"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(["view_use_legacy_sql"])}, client, None + {"path": "/%s" % path, "fields": json.dumps(["view_use_legacy_sql"])}, + client, + None, ) conn.api_request.assert_called_once_with( @@ -2676,7 +2692,9 @@ def test_update_table_w_query(self): updated_table = client.update_table(table, updated_properties) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(updated_properties)}, client, None + {"path": "/%s" % path, "fields": json.dumps(updated_properties)}, + client, + None, ) self.assertEqual(updated_table.schema, table.schema) @@ -2731,6 +2749,7 @@ def test_update_table_w_schema_None(self): self.TABLE_REF.table_id, ) ) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) table.schema = None @@ -2779,7 +2798,10 @@ def test_update_table_delete_property(self): table2 = client.update_table(table, ["description", "friendly_name"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(["description", "friendly_name"])}, + { + "path": "/%s" % path, + "fields": json.dumps(["description", "friendly_name"]), + }, client, None, ) @@ -3268,6 +3290,7 @@ def test_delete_model_w_not_found_ok_true(self): client.delete_model( "{}.{}".format(self.DS_ID, self.MODEL_ID), not_found_ok=True ) + final_attributes.assert_called_once_with( {"path": path, "not_found_okay": True}, client, None ) @@ -3346,6 +3369,7 @@ def test_delete_routine_w_not_found_ok_true(self): client.delete_routine( "routines-project.test_routines.test_routine", not_found_ok=True ) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with( From 5a32ad4b9d31c4ced636a52c6f8cb10ca62cfdd8 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Wed, 19 Aug 2020 10:55:40 -0400 Subject: [PATCH 27/41] adding Shawn's suggested changes --- google/cloud/bigquery/client.py | 8 +-- .../cloud/bigquery/opentelemetry_tracing.py | 4 +- tests/unit/test_client.py | 60 ++++++++----------- tests/unit/test_opentelemetry_tracing.py | 2 - 4 files changed, 30 insertions(+), 44 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index df98e43ba..2aa5a34fb 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -842,7 +842,7 @@ def update_dataset(self, dataset, fields, retry=DEFAULT_RETRY, timeout=None): else: headers = None path = dataset.path - span_attributes = {"path": path, "fields": json.dumps(fields)} + span_attributes = {"path": path, "fields": fields} with create_span( name="BigQuery.updateDataset", attributes=span_attributes, client=self @@ -891,7 +891,7 @@ def update_model(self, model, fields, retry=DEFAULT_RETRY, timeout=None): else: headers = None path = model.path - span_attributes = {"path": path, "fields": json.dumps(fields)} + span_attributes = {"path": path, "fields": fields} with create_span( name="BigQuery.updateModel", attributes=span_attributes, client=self @@ -951,7 +951,7 @@ def update_routine(self, routine, fields, retry=DEFAULT_RETRY, timeout=None): partial["routineReference"] = routine.reference.to_api_repr() path = routine.path - span_attributes = {"path": path, "fields": json.dumps(fields)} + span_attributes = {"path": path, "fields": fields} with create_span( name="BigQuery.updateRoutine", attributes=span_attributes, client=self @@ -1001,7 +1001,7 @@ def update_table(self, table, fields, retry=DEFAULT_RETRY, timeout=None): headers = None path = table.path - span_attributes = {"path": path, "fields": json.dumps(fields)} + span_attributes = {"path": path, "fields": fields} with create_span( name="BigQuery.updateTable", attributes=span_attributes, client=self diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 90ab99666..a7779eb0f 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -27,9 +27,9 @@ except ImportError: Logger.info( "This service is instrumented using OpenTelemetry." - "Opentelemetry could not be imported; please" + "OpenTelemetry could not be imported; please" "add opentelemetry-api and opentelemetry-instrumentation" - "packages in order to get Big Query Tracing data." + "packages in order to get BigQuery Tracing data." ) HAS_OPENTELEMETRY = False diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index dd82f2075..f3cbf0ce4 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -416,6 +416,7 @@ 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) projects = list(page) @@ -445,7 +446,6 @@ def test_list_projects_w_timeout(self): conn = client._connection = make_connection(DATA) iterator = client.list_projects(timeout=7.5) - six.next(iterator.pages) conn.api_request.assert_called_once_with( @@ -507,8 +507,8 @@ def test_list_datasets_defaults(self): creds = _make_credentials() client = self._make_one(self.PROJECT, creds) conn = client._connection = make_connection(DATA) - iterator = client.list_datasets() + iterator = client.list_datasets() page = six.next(iterator.pages) datasets = list(page) token = iterator.next_page_token @@ -550,7 +550,6 @@ 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) datasets = list(page) token = iterator.next_page_token @@ -1938,6 +1937,7 @@ def test_get_table_sets_user_agent(self): client_info=user_agent_override, _http=http, ) + client.get_table(self.TABLE_REF) expected_user_agent = user_agent_override.to_user_agent() @@ -2249,7 +2249,7 @@ def test_update_dataset(self): ds2 = client.update_dataset(ds, fields=fields, timeout=7.5,) final_attributes.assert_called_once_with( - {"path": "/" + PATH, "fields": json.dumps(fields)}, client, None + {"path": "/" + PATH, "fields": fields}, client, None ) conn.api_request.assert_called_once_with( @@ -2299,7 +2299,7 @@ def test_update_dataset_w_custom_property(self): dataset = client.update_dataset(dataset, ["newAlphaProperty"]) final_attributes.assert_called_once_with( - {"path": path, "fields": json.dumps(["newAlphaProperty"])}, client, None + {"path": path, "fields": ["newAlphaProperty"]}, client, None ) conn.api_request.assert_called_once_with( @@ -2355,7 +2355,7 @@ def test_update_model(self): updated_model = client.update_model(model, fields, timeout=7.5) final_attributes.assert_called_once_with( - {"path": "/" + path, "fields": json.dumps(fields)}, client, None + {"path": "/" + path, "fields": fields}, client, None ) sent = { @@ -2428,7 +2428,7 @@ def test_update_routine(self): actual_routine = client.update_routine(routine, fields, timeout=7.5,) final_attributes.assert_called_once_with( - {"path": routine.path, "fields": json.dumps(fields)}, client, None + {"path": routine.path, "fields": fields}, client, None ) # TODO: routineReference isn't needed when the Routines API supports @@ -2454,7 +2454,7 @@ def test_update_routine(self): client.update_routine(routine, []) final_attributes.assert_called_once_with( - {"path": routine.path, "fields": json.dumps([])}, client, None + {"path": routine.path, "fields": []}, client, None ) req = conn.api_request.call_args @@ -2515,7 +2515,7 @@ def test_update_table(self): span_path = "/%s" % path final_attributes.assert_called_once_with( - {"path": span_path, "fields": json.dumps(fields)}, client, None + {"path": span_path, "fields": fields}, client, None ) sent = { @@ -2555,7 +2555,7 @@ def test_update_table(self): client.update_table(table, []) final_attributes.assert_called_once_with( - {"path": "/" + path, "fields": json.dumps([])}, client, None + {"path": "/" + path, "fields": []}, client, None ) req = conn.api_request.call_args @@ -2583,9 +2583,7 @@ def test_update_table_w_custom_property(self): updated_table = client.update_table(table, ["newAlphaProperty"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(["newAlphaProperty"])}, - client, - None, + {"path": "/%s" % path, "fields": ["newAlphaProperty"]}, client, None, ) conn.api_request.assert_called_once_with( @@ -2620,9 +2618,7 @@ def test_update_table_only_use_legacy_sql(self): updated_table = client.update_table(table, ["view_use_legacy_sql"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(["view_use_legacy_sql"])}, - client, - None, + {"path": "/%s" % path, "fields": ["view_use_legacy_sql"]}, client, None, ) conn.api_request.assert_called_once_with( @@ -2692,9 +2688,7 @@ def test_update_table_w_query(self): updated_table = client.update_table(table, updated_properties) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(updated_properties)}, - client, - None, + {"path": "/%s" % path, "fields": updated_properties}, client, None, ) self.assertEqual(updated_table.schema, table.schema) @@ -2760,7 +2754,7 @@ def test_update_table_w_schema_None(self): updated_table = client.update_table(table, ["schema"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(["schema"])}, client, None + {"path": "/%s" % path, "fields": ["schema"]}, client, None ) self.assertEqual(len(conn.api_request.call_args_list), 2) @@ -2798,10 +2792,7 @@ def test_update_table_delete_property(self): table2 = client.update_table(table, ["description", "friendly_name"]) final_attributes.assert_called_once_with( - { - "path": "/%s" % path, - "fields": json.dumps(["description", "friendly_name"]), - }, + {"path": "/%s" % path, "fields": ["description", "friendly_name"],}, client, None, ) @@ -2815,7 +2806,7 @@ def test_update_table_delete_property(self): table3 = client.update_table(table2, ["description"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": json.dumps(["description"])}, client, None + {"path": "/%s" % path, "fields": ["description"]}, client, None ) self.assertEqual(len(conn.api_request.call_args_list), 2) @@ -2831,10 +2822,9 @@ def test_list_tables_empty_w_timeout(self): creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection({}) - dataset = DatasetReference(self.PROJECT, self.DS_ID) + 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) tables = list(page) @@ -2853,9 +2843,7 @@ def test_list_models_empty_w_timeout(self): conn = client._connection = make_connection({}) dataset_id = "{}.{}".format(self.PROJECT, self.DS_ID) - iterator = client.list_models(dataset_id, timeout=7.5) - page = six.next(iterator.pages) models = list(page) token = iterator.next_page_token @@ -2899,7 +2887,6 @@ def test_list_models_defaults(self): dataset = DatasetReference(self.PROJECT, self.DS_ID) iterator = client.list_models(dataset) - self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) models = list(page) @@ -2925,9 +2912,7 @@ def test_list_routines_empty_w_timeout(self): creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection({}) - path = "/projects/test-routines/datasets/test_routines/routines" iterator = client.list_routines("test-routines.test_routines", timeout=7.5) - page = six.next(iterator.pages) routines = list(page) token = iterator.next_page_token @@ -2935,7 +2920,10 @@ def test_list_routines_empty_w_timeout(self): self.assertEqual(routines, []) self.assertIsNone(token) conn.api_request.assert_called_once_with( - method="GET", path=path, query_params={}, timeout=7.5, + method="GET", + path="/projects/test-routines/datasets/test_routines/routines", + query_params={}, + timeout=7.5, ) def test_list_routines_defaults(self): @@ -2971,8 +2959,8 @@ def test_list_routines_defaults(self): client = self._make_one(project=project_id, credentials=creds) conn = client._connection = make_connection(resource) dataset = DatasetReference(client.project, dataset_id) - iterator = client.list_routines(dataset) + iterator = client.list_routines(dataset) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) routines = list(page) @@ -3035,8 +3023,8 @@ def test_list_tables_defaults(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) - iterator = client.list_tables(dataset) + iterator = client.list_tables(dataset) self.assertIs(iterator.dataset, dataset) page = six.next(iterator.pages) tables = list(page) @@ -3089,13 +3077,13 @@ def test_list_tables_explicit(self): client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection(DATA) dataset = DatasetReference(self.PROJECT, self.DS_ID) + iterator = client.list_tables( # Test with string for dataset ID. self.DS_ID, max_results=3, page_token=TOKEN, ) - self.assertEqual(iterator.dataset, dataset) page = six.next(iterator.pages) tables = list(page) diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 0064c9160..ec0d8834c 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -73,7 +73,6 @@ def test_opentelemetry_success(setup): @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_default_client_attributes(setup): - expected_attributes = { "foo": "baz", "db.system": "BigQuery", @@ -93,7 +92,6 @@ def test_default_client_attributes(setup): @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_default_job_attributes(setup): - expected_attributes = { "db.system": "BigQuery", "db.name": "test_project_id", From 54c133a64ef610db42bcebfe2df54f4f3a5b8e8c Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Wed, 19 Aug 2020 12:44:41 -0400 Subject: [PATCH 28/41] fixing _default_span_attribute_bug --- .../cloud/bigquery/opentelemetry_tracing.py | 16 +++---- setup.py | 1 + tests/unit/test_client.py | 1 - tests/unit/test_opentelemetry_tracing.py | 47 +++++++++++++++++++ 4 files changed, 56 insertions(+), 9 deletions(-) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index a7779eb0f..7f23164a0 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -14,6 +14,7 @@ import logging from contextlib import contextmanager +from frozendict import frozendict from google.api_core.exceptions import GoogleAPICallError Logger = logging.getLogger(__name__) @@ -34,9 +35,7 @@ HAS_OPENTELEMETRY = False -_default_attributes = { - "db.system": "BigQuery", -} +_default_attributes = frozendict({"db.system": "BigQuery",}) @contextmanager @@ -82,17 +81,18 @@ def create_span(name, attributes=None, client=None, job_ref=None): def _get_final_span_attributes(attributes=None, client=None, job_ref=None): + final_attributes = {} + final_attributes.update(_default_attributes.copy()) if client: client_attributes = _set_client_attributes(client) - _default_attributes.update(client_attributes) + final_attributes.update(client_attributes) if job_ref: job_attributes = _set_job_attributes(job_ref) - _default_attributes.update(job_attributes) + final_attributes.update(job_attributes) if attributes: - _default_attributes.update(attributes) - - return _default_attributes + final_attributes.update(attributes) + return final_attributes def _set_client_attributes(client): diff --git a/setup.py b/setup.py index 18bb78926..d23344652 100644 --- a/setup.py +++ b/setup.py @@ -62,6 +62,7 @@ "llvmlite<=0.34.0;python_version>='3.6'", "llvmlite<=0.31.0;python_version<'3.6'", ], + "frozendict": ["frozendict>=1.2"], } all_extras = [] diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index f3cbf0ce4..bba8c996d 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1937,7 +1937,6 @@ def test_get_table_sets_user_agent(self): client_info=user_agent_override, _http=http, ) - client.get_table(self.TABLE_REF) expected_user_agent = user_agent_override.to_user_agent() diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index ec0d8834c..9a78223f2 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -112,3 +112,50 @@ def test_default_job_attributes(setup): assert span is not None assert span.name == TEST_SPAN_NAME assert span.attributes == expected_attributes + + +@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") +def test_default_no_data_leakage(setup): + from google.cloud.bigquery import job + import google.auth.credentials + from google.cloud.bigquery import client + + mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) + test_client = client.Client( + project="test_project", credentials=mock_credentials, location="test_location" + ) + + expected_attributes = { + "foo": "baz", + "db.system": "BigQuery", + "db.name": "test_project", + "location": "test_location", + } + with opentelemetry_tracing.create_span( + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client + ) as span: + assert span.name == TEST_SPAN_NAME + assert span.attributes == expected_attributes + + test_job_reference = job._JobReference( + job_id="test_job_id", project="test_project_id", location="test_location" + ) + test_client = client.Client( + project="test_project", credentials=mock_credentials, location="test_location" + ) + test_job = job._AsyncJob(job_id=test_job_reference, client=test_client) + + expected_attributes = { + "db.system": "BigQuery", + "db.name": "test_project_id", + "location": "test_location", + "num_child_jobs": "0", + "job_id": "test_job_id", + "foo": "baz", + } + + with opentelemetry_tracing.create_span( + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job + ) as span: + assert span.name == TEST_SPAN_NAME + assert span.attributes == expected_attributes From 3aff14d9ee7d1129e73f482ab67dceea1a54f2ca Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Wed, 19 Aug 2020 13:21:12 -0400 Subject: [PATCH 29/41] reverting uneccesxsary changes --- google/cloud/bigquery/client.py | 19 +++++++++++-------- .../cloud/bigquery/opentelemetry_tracing.py | 2 +- tests/unit/test_client.py | 4 +++- tests/unit/test_job.py | 1 - 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 2aa5a34fb..e73930061 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -80,6 +80,7 @@ from google.cloud.bigquery.table import TableReference from google.cloud.bigquery.table import RowIterator + _DEFAULT_CHUNKSIZE = 1048576 # 1024 * 1024 B = 1 MB _MAX_MULTIPART_SIZE = 5 * 1024 * 1024 _DEFAULT_NUM_RETRIES = 6 @@ -352,11 +353,11 @@ def list_datasets( # TODO: consider supporting a dict of label -> value for filter, # and converting it into a string here. extra_params["filter"] = filter - + path = "/projects/%s/datasets" % (project,) return page_iterator.HTTPIterator( client=self, api_request=functools.partial(self._call_api, retry, timeout=timeout), - path="/projects/%s/datasets" % (project,), + path=path, item_to_value=_item_to_dataset, items_key="datasets", page_token=page_token, @@ -619,7 +620,6 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None): google.cloud.bigquery.dataset.Dataset: A ``Dataset`` instance. """ - if isinstance(dataset_ref, str): dataset_ref = DatasetReference.from_string( dataset_ref, default_project=self.project @@ -758,7 +758,6 @@ def get_routine(self, routine_ref, retry=DEFAULT_RETRY, timeout=None): google.cloud.bigquery.routine.Routine: A ``Routine`` instance. """ - if isinstance(routine_ref, str): routine_ref = RoutineReference.from_string( routine_ref, default_project=self.project @@ -1068,10 +1067,11 @@ def list_models( if not isinstance(dataset, (Dataset, DatasetReference)): raise TypeError("dataset must be a Dataset, DatasetReference, or string") + path = "%s/models" % dataset.path result = page_iterator.HTTPIterator( client=self, api_request=functools.partial(self._call_api, retry, timeout=timeout), - path="%s/models" % dataset.path, + path=path, item_to_value=_item_to_model, items_key="models", page_token=page_token, @@ -1132,10 +1132,11 @@ def list_routines( if not isinstance(dataset, (Dataset, DatasetReference)): raise TypeError("dataset must be a Dataset, DatasetReference, or string") + path = "{}/routines".format(dataset.path) result = page_iterator.HTTPIterator( client=self, api_request=functools.partial(self._call_api, retry, timeout=timeout), - path="{}/routines".format(dataset.path), + path=path, item_to_value=_item_to_routine, items_key="routines", page_token=page_token, @@ -1196,10 +1197,11 @@ def list_tables( if not isinstance(dataset, (Dataset, DatasetReference)): raise TypeError("dataset must be a Dataset, DatasetReference, or string") + path = "%s/tables" % dataset.path result = page_iterator.HTTPIterator( client=self, api_request=functools.partial(self._call_api, retry, timeout=timeout), - path="%s/tables" % dataset.path, + path=path, item_to_value=_item_to_table, items_key="tables", page_token=page_token, @@ -1763,10 +1765,11 @@ def list_jobs( if project is None: project = self.project + path = "/projects/%s/jobs" % (project,) return page_iterator.HTTPIterator( client=self, api_request=functools.partial(self._call_api, retry, timeout=timeout), - path="/projects/%s/jobs" % (project,), + path=path, item_to_value=_item_to_job, items_key="jobs", page_token=page_token, diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 7f23164a0..0bab0d7df 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -35,7 +35,7 @@ HAS_OPENTELEMETRY = False -_default_attributes = frozendict({"db.system": "BigQuery",}) +_default_attributes = frozendict({"db.system": "BigQuery"}) @contextmanager diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index bba8c996d..3ffe70d33 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1937,6 +1937,7 @@ def test_get_table_sets_user_agent(self): client_info=user_agent_override, _http=http, ) + client.get_table(self.TABLE_REF) expected_user_agent = user_agent_override.to_user_agent() @@ -2791,7 +2792,7 @@ def test_update_table_delete_property(self): table2 = client.update_table(table, ["description", "friendly_name"]) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "fields": ["description", "friendly_name"],}, + {"path": "/%s" % path, "fields": ["description", "friendly_name"]}, client, None, ) @@ -2911,6 +2912,7 @@ def test_list_routines_empty_w_timeout(self): creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) conn = client._connection = make_connection({}) + iterator = client.list_routines("test-routines.test_routines", timeout=7.5) page = six.next(iterator.pages) routines = list(page) diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 06b280d9b..d5497ffa8 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -2207,7 +2207,6 @@ def test_result_invokes_begin(self): client._connection = connection job = self._make_one(self.JOB_ID, [self.SOURCE1], self.TABLE_REF, client) - job.result() self.assertEqual(len(connection.api_request.call_args_list), 2) From 59477ff4c0d155f154e1b858f1217cb7fc9cf065 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Wed, 19 Aug 2020 16:10:39 -0400 Subject: [PATCH 30/41] adding more tests for all job_ref parameters --- .../cloud/bigquery/opentelemetry_tracing.py | 27 ++++++++++---- tests/unit/test_opentelemetry_tracing.py | 36 +++++++++++++++++-- 2 files changed, 53 insertions(+), 10 deletions(-) diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 0bab0d7df..3fafc076e 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -15,6 +15,7 @@ import logging from contextlib import contextmanager from frozendict import frozendict +import json from google.api_core.exceptions import GoogleAPICallError Logger = logging.getLogger(__name__) @@ -100,16 +101,28 @@ def _set_client_attributes(client): def _set_job_attributes(job_ref): - return { + job_attributes = { "db.name": job_ref.project, "location": job_ref.location, - "num_child_jobs": str(job_ref.num_child_jobs), + "num_child_jobs": job_ref.num_child_jobs, "job_id": job_ref.job_id, "parent_job_id": job_ref.parent_job_id, - "timeCreated": job_ref.created, - "timeStarted": job_ref.started, - "timeEnded": job_ref.ended, - "errors": job_ref.errors, - "errorResult": job_ref.error_result, "state": job_ref.state, } + if job_ref.errors is not None: + job_attributes["errors"] = json.dumps(job_ref.errors) + + if job_ref.error_result is not None: + job_attributes["errorResult"] = json.dumps(job_ref.error_result) + + date_time_format = "%d-%b-%Y (%H:%M:%S.%f)" + if job_ref.created is not None: + job_attributes["timeCreated"] = job_ref.created.strftime(date_time_format) + + if job_ref.started is not None: + job_attributes["timeStarted"] = job_ref.started.strftime(date_time_format) + + if job_ref.ended is not None: + job_attributes["timeEnded"] = job_ref.ended.strftime(date_time_format) + + return job_attributes diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 9a78223f2..4588f47d0 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -15,9 +15,9 @@ import sys import pytest - +import datetime import mock - +import json from google.cloud.bigquery import opentelemetry_tracing try: @@ -92,6 +92,22 @@ def test_default_client_attributes(setup): @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_default_job_attributes(setup): + import google.cloud._helpers + + time_created = datetime.datetime( + 2010, 5, 19, 16, 0, 0, tzinfo=google.cloud._helpers.UTC + ) + started_time = datetime.datetime( + 2011, 10, 1, 16, 0, 0, tzinfo=google.cloud._helpers.UTC + ) + ended_time = datetime.datetime( + 2011, 10, 2, 16, 0, 0, tzinfo=google.cloud._helpers.UTC + ) + errors = [{"error": "some_error"}] + error_result = [ + {"errorResult1": "some_error_result1", "errorResult2": "some_error_result2"} + ] + expected_attributes = { "db.system": "BigQuery", "db.name": "test_project_id", @@ -99,12 +115,26 @@ def test_default_job_attributes(setup): "num_child_jobs": "0", "job_id": "test_job_id", "foo": "baz", + "parent_job_id": "parent_job_id", + "timeCreated": time_created.strftime("%d-%b-%Y (%H:%M:%S.%f)"), + "timeStarted": started_time.strftime("%d-%b-%Y (%H:%M:%S.%f)"), + "timeEnded": ended_time.strftime("%d-%b-%Y (%H:%M:%S.%f)"), + "errors": json.dumps(errors), + "errorResult": json.dumps(error_result), + "state": "some_job_state", } with mock.patch("google.cloud.bigquery.job._AsyncJob") as test_job_ref: test_job_ref.job_id = "test_job_id" test_job_ref.location = "test_location" test_job_ref.project = "test_project_id" test_job_ref.num_child_jobs = "0" + test_job_ref.parent_job_id = "parent_job_id" + test_job_ref.created = time_created + test_job_ref.started = started_time + test_job_ref.ended = ended_time + test_job_ref.errors = errors + test_job_ref.error_result = error_result + test_job_ref.state = "some_job_state" with opentelemetry_tracing.create_span( TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job_ref @@ -149,7 +179,7 @@ def test_default_no_data_leakage(setup): "db.system": "BigQuery", "db.name": "test_project_id", "location": "test_location", - "num_child_jobs": "0", + "num_child_jobs": 0, "job_id": "test_job_id", "foo": "baz", } From 0d15998a5688aa74c72fa77c14ef4441ea139a5d Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Wed, 19 Aug 2020 23:47:37 -0400 Subject: [PATCH 31/41] removing dependecny, ordering imports and other changes --- google/cloud/bigquery/client.py | 2 +- google/cloud/bigquery/job.py | 2 +- .../cloud/bigquery/opentelemetry_tracing.py | 16 +++---- setup.py | 1 - tests/unit/test_client.py | 17 +++---- tests/unit/test_opentelemetry_tracing.py | 44 ++++++++++++++----- 6 files changed, 49 insertions(+), 33 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index e73930061..5c50492ba 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -33,7 +33,6 @@ import tempfile import uuid import warnings -from google.cloud.bigquery.opentelemetry_tracing import create_span try: import pyarrow @@ -64,6 +63,7 @@ from google.cloud.bigquery.dataset import DatasetListItem from google.cloud.bigquery.dataset import DatasetReference from google.cloud.bigquery.exceptions import PyarrowMissingWarning +from google.cloud.bigquery.opentelemetry_tracing import create_span from google.cloud.bigquery import job from google.cloud.bigquery.model import Model from google.cloud.bigquery.model import ModelReference diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index 6db71b551..e7ab38da9 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -34,6 +34,7 @@ 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 @@ -50,7 +51,6 @@ from google.cloud.bigquery.table import TableReference from google.cloud.bigquery.table import Table from google.cloud.bigquery.table import TimePartitioning -from google.cloud.bigquery.opentelemetry_tracing import create_span _DONE_STATE = "DONE" diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 3fafc076e..f1ad9c5d0 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -14,15 +14,13 @@ import logging from contextlib import contextmanager -from frozendict import frozendict -import json from google.api_core.exceptions import GoogleAPICallError Logger = logging.getLogger(__name__) try: from opentelemetry import trace - from opentelemetry.trace.status import Status from opentelemetry.instrumentation.utils import http_status_to_canonical_code + from opentelemetry.trace.status import Status HAS_OPENTELEMETRY = True @@ -36,7 +34,9 @@ HAS_OPENTELEMETRY = False -_default_attributes = frozendict({"db.system": "BigQuery"}) +_default_attributes = { + "db.system": "BigQuery" +} # static, default values assigned to all spans @contextmanager @@ -74,7 +74,6 @@ def create_span(name, attributes=None, client=None, job_ref=None): with tracer.start_as_current_span(name=name, attributes=final_attributes) as span: try: yield span - span.set_status(Status(http_status_to_canonical_code(200))) except GoogleAPICallError as error: if error.code is not None: span.set_status(Status(http_status_to_canonical_code(error.code))) @@ -109,11 +108,12 @@ def _set_job_attributes(job_ref): "parent_job_id": job_ref.parent_job_id, "state": job_ref.state, } - if job_ref.errors is not None: - job_attributes["errors"] = json.dumps(job_ref.errors) if job_ref.error_result is not None: - job_attributes["errorResult"] = json.dumps(job_ref.error_result) + job_attributes["hasErrors"] = True + + else: + job_attributes["hasErrors"] = False date_time_format = "%d-%b-%Y (%H:%M:%S.%f)" if job_ref.created is not None: diff --git a/setup.py b/setup.py index d23344652..18bb78926 100644 --- a/setup.py +++ b/setup.py @@ -62,7 +62,6 @@ "llvmlite<=0.34.0;python_version>='3.6'", "llvmlite<=0.31.0;python_version<'3.6'", ], - "frozendict": ["frozendict>=1.2"], } all_extras = [] diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 3ffe70d33..d64e981a3 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1135,14 +1135,7 @@ def test_create_dataset_alreadyexists_w_exists_ok_false(self): ) with pytest.raises(google.api_core.exceptions.AlreadyExists): - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - client.create_dataset(self.DS_ID) - - final_attributes.assert_called_once_with( - {"path": "path"}, client, None - ) # pragma: NO COVER + client.create_dataset(self.DS_ID) def test_create_dataset_alreadyexists_w_exists_ok_true(self): post_path = "/projects/{}/datasets".format(self.PROJECT) @@ -1281,7 +1274,7 @@ def test_create_routine_w_conflict_exists_ok(self): actual_routine = client.create_routine(routine, exists_ok=True) final_attributes.assert_called_with( - {"path": path + "/minimal_routine"}, client, None + {"path": "%s/minimal_routine" % path}, client, None ) self.assertEqual(actual_routine.project, "test-routine-project") @@ -2249,7 +2242,7 @@ def test_update_dataset(self): ds2 = client.update_dataset(ds, fields=fields, timeout=7.5,) final_attributes.assert_called_once_with( - {"path": "/" + PATH, "fields": fields}, client, None + {"path": "/%s" % PATH, "fields": fields}, client, None ) conn.api_request.assert_called_once_with( @@ -2355,7 +2348,7 @@ def test_update_model(self): updated_model = client.update_model(model, fields, timeout=7.5) final_attributes.assert_called_once_with( - {"path": "/" + path, "fields": fields}, client, None + {"path": "/%s" % path, "fields": fields}, client, None ) sent = { @@ -2555,7 +2548,7 @@ def test_update_table(self): client.update_table(table, []) final_attributes.assert_called_once_with( - {"path": "/" + path, "fields": []}, client, None + {"path": "/%s" % path, "fields": []}, client, None ) req = conn.api_request.call_args diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 4588f47d0..a2646a70c 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -12,13 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. - -import sys -import pytest import datetime import mock -import json +import pytest +import sys +import unittest + from google.cloud.bigquery import opentelemetry_tracing +from six.moves import reload_module try: import opentelemetry @@ -31,7 +32,6 @@ except ImportError: opentelemetry = None -from six.moves import reload_module TEST_SPAN_NAME = "bar" TEST_SPAN_ATTRIBUTES = {"foo": "baz"} @@ -103,7 +103,6 @@ def test_default_job_attributes(setup): ended_time = datetime.datetime( 2011, 10, 2, 16, 0, 0, tzinfo=google.cloud._helpers.UTC ) - errors = [{"error": "some_error"}] error_result = [ {"errorResult1": "some_error_result1", "errorResult2": "some_error_result2"} ] @@ -119,8 +118,7 @@ def test_default_job_attributes(setup): "timeCreated": time_created.strftime("%d-%b-%Y (%H:%M:%S.%f)"), "timeStarted": started_time.strftime("%d-%b-%Y (%H:%M:%S.%f)"), "timeEnded": ended_time.strftime("%d-%b-%Y (%H:%M:%S.%f)"), - "errors": json.dumps(errors), - "errorResult": json.dumps(error_result), + "hasErrors": True, "state": "some_job_state", } with mock.patch("google.cloud.bigquery.job._AsyncJob") as test_job_ref: @@ -132,7 +130,6 @@ def test_default_job_attributes(setup): test_job_ref.created = time_created test_job_ref.started = started_time test_job_ref.ended = ended_time - test_job_ref.errors = errors test_job_ref.error_result = error_result test_job_ref.state = "some_job_state" @@ -146,9 +143,9 @@ def test_default_job_attributes(setup): @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") def test_default_no_data_leakage(setup): - from google.cloud.bigquery import job import google.auth.credentials from google.cloud.bigquery import client + from google.cloud.bigquery import job mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) test_client = client.Client( @@ -182,6 +179,7 @@ def test_default_no_data_leakage(setup): "num_child_jobs": 0, "job_id": "test_job_id", "foo": "baz", + "hasErrors": False, } with opentelemetry_tracing.create_span( @@ -189,3 +187,29 @@ def test_default_no_data_leakage(setup): ) as span: assert span.name == TEST_SPAN_NAME assert span.attributes == expected_attributes + + +@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") +def test_span_creation_error(setup): + import google.auth.credentials + from google.cloud.bigquery import client + from google.api_core.exceptions import GoogleAPICallError, InvalidArgument + + mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials) + test_client = client.Client( + project="test_project", credentials=mock_credentials, location="test_location" + ) + + expected_attributes = { + "foo": "baz", + "db.system": "BigQuery", + "db.name": "test_project", + "location": "test_location", + } + with unittest.TestCase().assertRaises(GoogleAPICallError): + with opentelemetry_tracing.create_span( + TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client + ) as span: + assert span.name == TEST_SPAN_NAME + assert span.attributes == expected_attributes + raise InvalidArgument("test_error") From d9971752c8b2006210f28efcb203300489224914 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Thu, 20 Aug 2020 10:06:39 -0400 Subject: [PATCH 32/41] addressing Shawn concerns --- google/cloud/bigquery/client.py | 1 - google/cloud/bigquery/job.py | 1 - google/cloud/bigquery/opentelemetry_tracing.py | 18 ++++++------------ tests/unit/test_client.py | 1 + tests/unit/test_opentelemetry_tracing.py | 6 +++--- 5 files changed, 10 insertions(+), 17 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 5c50492ba..5a27cbc27 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -289,7 +289,6 @@ def list_projects( Iterator of :class:`~google.cloud.bigquery.client.Project` accessible to the current client. """ - return page_iterator.HTTPIterator( client=self, api_request=functools.partial(self._call_api, retry, timeout=timeout), diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index e7ab38da9..a8e0c25ed 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -52,7 +52,6 @@ from google.cloud.bigquery.table import Table from google.cloud.bigquery.table import TimePartitioning - _DONE_STATE = "DONE" _STOPPED_REASON = "stopped" _TIMEOUT_BUFFER_SECS = 0.1 diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index f1ad9c5d0..6eb8833f4 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -41,8 +41,8 @@ @contextmanager def create_span(name, attributes=None, client=None, job_ref=None): - """Creates a ContextManager for a Span to be exported to the configured exporter. If no configuration - exists yields None. + """Creates a ContextManager for a Span to be exported to the configured exporter. + If no configuration exists yields None. Args: name (str): Name that will be set for the span being created @@ -89,7 +89,6 @@ def _get_final_span_attributes(attributes=None, client=None, job_ref=None): if job_ref: job_attributes = _set_job_attributes(job_ref) final_attributes.update(job_attributes) - if attributes: final_attributes.update(attributes) return final_attributes @@ -109,20 +108,15 @@ def _set_job_attributes(job_ref): "state": job_ref.state, } - if job_ref.error_result is not None: - job_attributes["hasErrors"] = True - - else: - job_attributes["hasErrors"] = False + job_attributes["hasErrors"] = job_ref.error_result is not None - date_time_format = "%d-%b-%Y (%H:%M:%S.%f)" if job_ref.created is not None: - job_attributes["timeCreated"] = job_ref.created.strftime(date_time_format) + job_attributes["timeCreated"] = job_ref.created.isoformat() if job_ref.started is not None: - job_attributes["timeStarted"] = job_ref.started.strftime(date_time_format) + job_attributes["timeStarted"] = job_ref.started.isoformat() if job_ref.ended is not None: - job_attributes["timeEnded"] = job_ref.ended.strftime(date_time_format) + job_attributes["timeEnded"] = job_ref.ended.isoformat() return job_attributes diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index d64e981a3..4b53b666f 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -89,6 +89,7 @@ def _make_list_partitons_meta_info(project, dataset_id, table_id, num_rows=0): class TestClient(unittest.TestCase): + PROJECT = "PROJECT" DS_ID = "DATASET_ID" TABLE_ID = "TABLE_ID" diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index a2646a70c..353ff328f 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -115,9 +115,9 @@ def test_default_job_attributes(setup): "job_id": "test_job_id", "foo": "baz", "parent_job_id": "parent_job_id", - "timeCreated": time_created.strftime("%d-%b-%Y (%H:%M:%S.%f)"), - "timeStarted": started_time.strftime("%d-%b-%Y (%H:%M:%S.%f)"), - "timeEnded": ended_time.strftime("%d-%b-%Y (%H:%M:%S.%f)"), + "timeCreated": time_created.isoformat(), + "timeStarted": started_time.isoformat(), + "timeEnded": ended_time.isoformat(), "hasErrors": True, "state": "some_job_state", } From 61b5b858f67c01464ee9d391367452d155fa665c Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Thu, 20 Aug 2020 12:38:26 -0400 Subject: [PATCH 33/41] adding test and suggested changes --- README.rst | 4 +- google/cloud/bigquery/client.py | 2 +- .../cloud/bigquery/opentelemetry_tracing.py | 4 +- tests/unit/test_client.py | 94 +++++++++++-------- 4 files changed, 62 insertions(+), 42 deletions(-) diff --git a/README.rst b/README.rst index 4cfde7076..0457f5177 100644 --- a/README.rst +++ b/README.rst @@ -118,8 +118,8 @@ the BigQuery client the following PyPI packages need to be installed: After installation, OpenTelemetry can be used in the BigQuery client and in BigQuery jobs. First, however, an exporter must be -specified for where the trace data will be outputted to. An example of this -can be found here: +specified for where the trace data will be outputted to. An +example of this can be found here: .. code-block:: python diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 5a27cbc27..38dec325a 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -1311,7 +1311,7 @@ def delete_model( path = model.path try: - span_attributes = {"not_found_okay": not_found_ok, "path": path} + span_attributes = {"path": path} with create_span( name="BigQuery.deleteModel", attributes=span_attributes, client=self ): diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py index 6eb8833f4..f7375c346 100644 --- a/google/cloud/bigquery/opentelemetry_tracing.py +++ b/google/cloud/bigquery/opentelemetry_tracing.py @@ -16,7 +16,7 @@ from contextlib import contextmanager from google.api_core.exceptions import GoogleAPICallError -Logger = logging.getLogger(__name__) +logger = logging.getLogger(__name__) try: from opentelemetry import trace from opentelemetry.instrumentation.utils import http_status_to_canonical_code @@ -25,7 +25,7 @@ HAS_OPENTELEMETRY = True except ImportError: - Logger.info( + logger.info( "This service is instrumented using OpenTelemetry." "OpenTelemetry could not be imported; please" "add opentelemetry-api and opentelemetry-instrumentation" diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 4b53b666f..2aad14b78 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -40,6 +40,16 @@ import pandas except (ImportError, AttributeError): # pragma: NO COVER pandas = None +try: + import opentelemetry + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) +except (ImportError, AttributeError): + opentelemetry = None try: import pyarrow except (ImportError, AttributeError): # pragma: NO COVER @@ -651,9 +661,7 @@ def test_get_dataset(self): ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with( # pragma: NO COVER - {"path": "/%s" % path}, client, None - ) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) # Zero-length errors field. client._connection = make_connection(ServerError(""), resource) @@ -663,9 +671,7 @@ def test_get_dataset(self): ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with( # pragma: NO COVER - {"path": "/%s" % path}, client, None - ) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) # Non-retryable reason. client._connection = make_connection( @@ -691,9 +697,7 @@ def test_get_dataset(self): ) as final_attributes: client.get_dataset(dataset_ref, retry=None) - final_attributes.assert_called_once_with( # pragma: NO COVER - {"path": "/%s" % path}, client, None - ) + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) # Retryable reason, default retry: success. client._connection = make_connection( @@ -1235,9 +1239,43 @@ def test_create_routine_w_conflict(self): ) as final_attributes: client.create_routine(routine) - final_attributes.assert_called_once_with( # pragma: NO COVER - {"path": path, "exists_okay": False}, client, None - ) + final_attributes.assert_called_once_with({"path": path}, client, None) + + resource = { + "routineReference": { + "projectId": "test-routine-project", + "datasetId": "test_routines", + "routineId": "minimal_routine", + } + } + conn.api_request.assert_called_once_with( + method="POST", path=path, data=resource, timeout=None, + ) + + @unittest.skipIf(opentelemetry is None, "Requires `opentelemetry`") + def test_span_status_is_set(self): + from google.cloud.bigquery.routine import Routine + + tracer_provider = TracerProvider() + memory_exporter = InMemorySpanExporter() + span_processor = SimpleExportSpanProcessor(memory_exporter) + tracer_provider.add_span_processor(span_processor) + trace.set_tracer_provider(tracer_provider) + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + conn = client._connection = make_connection( + google.api_core.exceptions.AlreadyExists("routine already exists") + ) + path = "/projects/test-routine-project/datasets/test_routines/routines" + full_routine_id = "test-routine-project.test_routines.minimal_routine" + routine = Routine(full_routine_id) + + with pytest.raises(google.api_core.exceptions.AlreadyExists): + client.create_routine(routine) + + span_list = memory_exporter.get_finished_spans() + self.assertTrue(span_list[0].status is not None) resource = { "routineReference": { @@ -1701,11 +1739,9 @@ def test_create_table_alreadyexists_w_exists_ok_false(self): ) as final_attributes: client.create_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_with( # pragma: NO COVER - {"path": post_path, "dataset_id": self.TABLE_REF.dataset_id}, - client, - None, - ) + final_attributes.assert_called_with( + {"path": post_path, "dataset_id": self.TABLE_REF.dataset_id}, client, None, + ) conn.api_request.assert_called_once_with( method="POST", @@ -2113,14 +2149,7 @@ def test_set_iam_policy_invalid_policy(self): client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) with self.assertRaises(TypeError): - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - client.set_iam_policy(self.TABLE_REF, invalid_policy_repr) - - final_attributes.assert_called_once_with( # pragma: NO COVER - {"path": "{}:setIamPolicy".format(self.TABLE_REF.path)}, client, None - ) + client.set_iam_policy(self.TABLE_REF, invalid_policy_repr) def test_set_iam_policy_w_invalid_table(self): from google.api_core.iam import Policy @@ -2138,14 +2167,7 @@ def test_set_iam_policy_w_invalid_table(self): ) with self.assertRaises(TypeError): - with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" - ) as final_attributes: - client.set_iam_policy(table_resource_string, policy) - - final_attributes.assert_called_once_with( # pragma: NO COVER - {"path": table_resource_string}, client, None - ) + client.set_iam_policy(table_resource_string, policy) def test_test_iam_permissions(self): PATH = "/projects/%s/datasets/%s/tables/%s:testIamPermissions" % ( @@ -3229,7 +3251,7 @@ def test_delete_model(self): client.delete_model(arg, timeout=7.5) final_attributes.assert_called_once_with( - {"path": "/%s" % path, "not_found_okay": False}, client, None + {"path": "/%s" % path}, client, None ) conn.api_request.assert_called_with( method="DELETE", path="/%s" % path, timeout=7.5 @@ -3274,9 +3296,7 @@ def test_delete_model_w_not_found_ok_true(self): "{}.{}".format(self.DS_ID, self.MODEL_ID), not_found_ok=True ) - final_attributes.assert_called_once_with( - {"path": path, "not_found_okay": True}, client, None - ) + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) From 4d1dec11a72c19ea743ec6ca855551eafb223a47 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Thu, 20 Aug 2020 15:10:07 -0400 Subject: [PATCH 34/41] adding opentelemetry to setup.py and other suggested changes --- README.rst | 2 +- google/cloud/bigquery/client.py | 2 +- noxfile.py | 7 ++----- setup.py | 7 ++++++- tests/unit/test_client.py | 20 +++++--------------- 5 files changed, 15 insertions(+), 23 deletions(-) diff --git a/README.rst b/README.rst index 0457f5177..c6bc17834 100644 --- a/README.rst +++ b/README.rst @@ -114,7 +114,7 @@ the BigQuery client the following PyPI packages need to be installed: .. code-block:: console - pip install opentelemetry-api opentelemetry-sdk opentelemetry-instrumentation opentelemetry-exporter-google-cloud + pip install google-cloud-bigquery[opentelemetry] opentelemetry-exporter-google-cloud After installation, OpenTelemetry can be used in the BigQuery client and in BigQuery jobs. First, however, an exporter must be diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 38dec325a..fbbfda051 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -1451,7 +1451,7 @@ def _get_query_results( # This call is typically made in a polling loop that checks whether the # job is complete (from QueryJob.done(), called ultimately from # QueryJob.result()). So we don't need to poll here. - span_attributes = {"path": path, "job_id": job_id, "location": location} + span_attributes = {"path": path} with create_span( name="BigQuery.getQueryResults", attributes=span_attributes, client=self diff --git a/noxfile.py b/noxfile.py index e4aad1866..5db14c31f 100644 --- a/noxfile.py +++ b/noxfile.py @@ -51,11 +51,8 @@ def default(session): # opentelemetry was not added to [all] because opentelemetry does not support Python 2. # Exporter does not need to be in nox thus it has been added to README documentation if session.python != "2.7": - session.install( - "opentelemetry-api==0.9b0", - "opentelemetry-sdk==0.9b0", - "opentelemetry-instrumentation==0.9b0 ", - ) + session.install("-e", ".[opentelemetry]") + # Run py.test against the unit tests. session.run( "py.test", diff --git a/setup.py b/setup.py index 18bb78926..22f4e8500 100644 --- a/setup.py +++ b/setup.py @@ -62,12 +62,17 @@ "llvmlite<=0.34.0;python_version>='3.6'", "llvmlite<=0.31.0;python_version<'3.6'", ], + "opentelemetry": [ + "opentelemetry-api==0.9b0", + "opentelemetry-sdk==0.9b0", + "opentelemetry-instrumentation==0.9b0 ", + ], } all_extras = [] for extra in extras: - if extra == "fastparquet": + if extra == "fastparquet" or extra == "opentelemetry": # Skip fastparquet from "all" because it is redundant with pyarrow and # creates a dependency on pre-release versions of numpy. See: # https://github.com/googleapis/google-cloud-python/issues/8549 diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 2aad14b78..271640dd5 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -270,9 +270,7 @@ def test__get_query_results_miss_w_explicit_project_and_timeout(self): timeout=42, ) - final_attributes.assert_called_once_with( - {"path": path}, client, None - ) # pragma: NO COVER + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_once_with( method="GET", @@ -683,9 +681,7 @@ def test_get_dataset(self): ) as final_attributes: client.get_dataset(dataset_ref) - final_attributes.assert_called_once_with( - {"path": path}, client, None - ) # pragma: NO COVER + final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None) # Retryable reason, but retry is disabled. client._connection = make_connection( @@ -3197,9 +3193,7 @@ def test_delete_dataset_w_not_found_ok_false(self): ) as final_attributes: client.delete_dataset(self.DS_ID) - final_attributes.assert_called_once_with( - {"path": path}, client, None - ) # pragma: NO COVER + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with( method="DELETE", path=path, query_params={}, timeout=None @@ -3349,9 +3343,7 @@ def test_delete_routine_w_not_found_ok_false(self): ) as final_attributes: client.delete_routine("routines-project.test_routines.test_routine") - final_attributes.assert_called_once_with( - {"path": path}, client, None - ) # pragma: NO COVER + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with( method="DELETE", path=path, timeout=None, @@ -3438,9 +3430,7 @@ def test_delete_table_w_not_found_ok_false(self): ) as final_attributes: client.delete_table("{}.{}".format(self.DS_ID, self.TABLE_ID)) - final_attributes.assert_called_once_with( - {"path": path}, client, None - ) # pragma: NO COVER + final_attributes.assert_called_once_with({"path": path}, client, None) conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None) From 7704975ecd1ac93b49a9c5f55d87ddd1bc5746a8 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Thu, 20 Aug 2020 15:52:21 -0400 Subject: [PATCH 35/41] adding reasoning for not adding to [all] --- setup.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index 22f4e8500..916b5ec48 100644 --- a/setup.py +++ b/setup.py @@ -72,10 +72,15 @@ all_extras = [] for extra in extras: - if extra == "fastparquet" or extra == "opentelemetry": - # Skip fastparquet from "all" because it is redundant with pyarrow and - # creates a dependency on pre-release versions of numpy. See: - # https://github.com/googleapis/google-cloud-python/issues/8549 + if extra in ( + # Skip fastparquet from "all" because it is redundant with pyarrow and + # creates a dependency on pre-release versions of numpy. See: + # https://github.com/googleapis/google-cloud-python/issues/8549 + "fastparquet", + + # Skip opentelemetry because the library is not compatible with Python 2. + "opentelemetry", + ): continue all_extras.extend(extras[extra]) From 839cd92c872582bbf6f787cbaed28843d2c4e9f9 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Thu, 20 Aug 2020 16:36:00 -0400 Subject: [PATCH 36/41] linting --- setup.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/setup.py b/setup.py index 916b5ec48..f30968364 100644 --- a/setup.py +++ b/setup.py @@ -73,13 +73,12 @@ for extra in extras: if extra in ( - # Skip fastparquet from "all" because it is redundant with pyarrow and - # creates a dependency on pre-release versions of numpy. See: - # https://github.com/googleapis/google-cloud-python/issues/8549 - "fastparquet", - - # Skip opentelemetry because the library is not compatible with Python 2. - "opentelemetry", + # Skip fastparquet from "all" because it is redundant with pyarrow and + # creates a dependency on pre-release versions of numpy. See: + # https://github.com/googleapis/google-cloud-python/issues/8549 + "fastparquet", + # Skip opentelemetry because the library is not compatible with Python 2. + "opentelemetry", ): continue all_extras.extend(extras[extra]) From 55c322eb2b96f423e09f64a5a6101bd9bee7c2a9 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Fri, 21 Aug 2020 11:34:38 -0400 Subject: [PATCH 37/41] adding nested functions --- google/cloud/bigquery/client.py | 63 +++++++++++++++++++++++++++++---- 1 file changed, 57 insertions(+), 6 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index fbbfda051..e2f4896f2 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", From 4f66ffb2c39c7c5e450ee8652a0beda4e3ce30d7 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Fri, 21 Aug 2020 11:35:46 -0400 Subject: [PATCH 38/41] adding test --- tests/unit/test_client.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 271640dd5..5063fd296 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -425,8 +425,12 @@ def test_list_projects_defaults(self): creds = _make_credentials() client = self._make_one(PROJECT_1, creds) conn = client._connection = make_connection(DATA) + with mock.patch( + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + ) as final_attributes: + iterator = client.list_projects() - iterator = client.list_projects() + final_attributes.assert_called_once_with({"path": '/projects'}, client, None) page = six.next(iterator.pages) projects = list(page) token = iterator.next_page_token From dd1e2462feb8b3ff3f2d66615d8b3ce073a34532 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Fri, 21 Aug 2020 12:20:50 -0400 Subject: [PATCH 39/41] adding Tim suggested changes --- tests/unit/test_opentelemetry_tracing.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py index 353ff328f..1c35b0a82 100644 --- a/tests/unit/test_opentelemetry_tracing.py +++ b/tests/unit/test_opentelemetry_tracing.py @@ -13,13 +13,9 @@ # limitations under the License. import datetime -import mock -import pytest import sys -import unittest -from google.cloud.bigquery import opentelemetry_tracing -from six.moves import reload_module +import mock try: import opentelemetry @@ -31,7 +27,10 @@ ) except ImportError: opentelemetry = None +import pytest +from six.moves import reload_module +from google.cloud.bigquery import opentelemetry_tracing TEST_SPAN_NAME = "bar" TEST_SPAN_ATTRIBUTES = {"foo": "baz"} @@ -40,6 +39,7 @@ @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") @pytest.fixture def setup(): + reload_module(opentelemetry_tracing) tracer_provider = TracerProvider() memory_exporter = InMemorySpanExporter() span_processor = SimpleExportSpanProcessor(memory_exporter) @@ -49,14 +49,11 @@ def setup(): @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") -def test_opentelemetry_not_installed(setup): - temp_module = sys.modules["opentelemetry"] - sys.modules["opentelemetry"] = None +def test_opentelemetry_not_installed(setup, monkeypatch): + monkeypatch.setitem(sys.modules, "opentelemetry", None) reload_module(opentelemetry_tracing) with opentelemetry_tracing.create_span("No-op for opentelemetry") as span: assert span is None - sys.modules["opentelemetry"] = temp_module - reload_module(opentelemetry_tracing) @pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`") @@ -206,7 +203,7 @@ def test_span_creation_error(setup): "db.name": "test_project", "location": "test_location", } - with unittest.TestCase().assertRaises(GoogleAPICallError): + with pytest.raises(GoogleAPICallError): with opentelemetry_tracing.create_span( TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client ) as span: From 5b3a202e4807f8d9ec7dd19a2394189e4a0fe4a2 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Fri, 21 Aug 2020 23:32:12 -0400 Subject: [PATCH 40/41] adding full tests --- tests/unit/test_client.py | 110 ++++++++++++++++++++++++++++++++------ 1 file changed, 93 insertions(+), 17 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 5063fd296..65ee8a63c 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -425,13 +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() + with mock.patch( "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: - iterator = client.list_projects() + page = six.next(iterator.pages) final_attributes.assert_called_once_with({"path": '/projects'}, client, None) - page = six.next(iterator.pages) projects = list(page) token = iterator.next_page_token @@ -459,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 @@ -473,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 @@ -522,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 @@ -542,7 +560,12 @@ 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", @@ -563,7 +586,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 @@ -2842,7 +2870,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 @@ -2860,7 +2893,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 @@ -2904,7 +2942,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 @@ -2930,7 +2973,12 @@ 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 @@ -2979,7 +3027,10 @@ 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) routines = list(page) actual_token = iterator.next_page_token @@ -3043,7 +3094,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 @@ -3102,7 +3158,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 @@ -3925,7 +3986,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 @@ -3970,7 +4036,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 @@ -3999,7 +4070,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 From 9e087ff755a347072cc642fbbe34c79c33c18fc0 Mon Sep 17 00:00:00 2001 From: Aravin Sivakumar Date: Sun, 23 Aug 2020 23:36:45 -0400 Subject: [PATCH 41/41] linting and fixing missing test --- google/cloud/bigquery/client.py | 12 ++++++------ tests/unit/test_client.py | 24 ++++++++++++++++-------- 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index e2f4896f2..e7f8c5c66 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -289,11 +289,11 @@ def list_projects( Iterator of :class:`~google.cloud.bigquery.client.Project` accessible to the current client. """ - span_attributes = {"path": '/projects'} + span_attributes = {"path": "/projects"} def api_request(*args, **kwargs): with create_span( - name="BigQuery.listProjects", attributes=span_attributes, client=self + name="BigQuery.listProjects", attributes=span_attributes, client=self ): return self._call_api(retry, *args, timeout=timeout, **kwargs) @@ -366,7 +366,7 @@ def list_datasets( def api_request(*args, **kwargs): with create_span( - name="BigQuery.listDatasets", attributes=span_attributes, client=self + name="BigQuery.listDatasets", attributes=span_attributes, client=self ): return self._call_api(retry, *args, timeout=timeout, **kwargs) @@ -1162,7 +1162,7 @@ def list_routines( def api_request(*args, **kwargs): with create_span( - name="BigQuery.listRoutines", attributes=span_attributes, client=self + name="BigQuery.listRoutines", attributes=span_attributes, client=self ): return self._call_api(retry, *args, timeout=timeout, **kwargs) @@ -1235,7 +1235,7 @@ def list_tables( def api_request(*args, **kwargs): with create_span( - name="BigQuery.listTables", attributes=span_attributes, client=self + name="BigQuery.listTables", attributes=span_attributes, client=self ): return self._call_api(retry, *args, timeout=timeout, **kwargs) @@ -1812,7 +1812,7 @@ def list_jobs( def api_request(*args, **kwargs): with create_span( - name="BigQuery.listJobs", attributes=span_attributes, client=self + name="BigQuery.listJobs", attributes=span_attributes, client=self ): return self._call_api(retry, *args, timeout=timeout, **kwargs) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 65ee8a63c..01bb1f2e1 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -432,7 +432,7 @@ def test_list_projects_defaults(self): ) as final_attributes: page = six.next(iterator.pages) - final_attributes.assert_called_once_with({"path": '/projects'}, client, None) + final_attributes.assert_called_once_with({"path": "/projects"}, client, None) projects = list(page) token = iterator.next_page_token @@ -466,7 +466,7 @@ def test_list_projects_w_timeout(self): ) as final_attributes: six.next(iterator.pages) - final_attributes.assert_called_once_with({"path": '/projects'}, client, None) + 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 @@ -486,7 +486,7 @@ def test_list_projects_explicit_response_missing_projects_key(self): ) as final_attributes: page = six.next(iterator.pages) - final_attributes.assert_called_once_with({"path": '/projects'}, client, None) + final_attributes.assert_called_once_with({"path": "/projects"}, client, None) projects = list(page) token = iterator.next_page_token @@ -565,7 +565,9 @@ def test_list_datasets_w_project_and_timeout(self): ) 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) + final_attributes.assert_called_once_with( + {"path": "/projects/other-project/datasets"}, client, None + ) conn.api_request.assert_called_once_with( method="GET", @@ -2978,7 +2980,11 @@ def test_list_routines_empty_w_timeout(self): ) as final_attributes: page = six.next(iterator.pages) - final_attributes.assert_called_once_with({"path": "/projects/test-routines/datasets/test_routines/routines"}, client, None) + final_attributes.assert_called_once_with( + {"path": "/projects/test-routines/datasets/test_routines/routines"}, + client, + None, + ) routines = list(page) token = iterator.next_page_token @@ -3031,6 +3037,8 @@ def test_list_routines_defaults(self): "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 @@ -3159,7 +3167,7 @@ def test_list_tables_explicit(self): ) self.assertEqual(iterator.dataset, dataset) with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: page = six.next(iterator.pages) @@ -4037,7 +4045,7 @@ def test_list_jobs_load_job_wo_sourceUris(self): iterator = client.list_jobs() with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: page = six.next(iterator.pages) @@ -4071,7 +4079,7 @@ def test_list_jobs_explicit_missing(self): max_results=1000, page_token=TOKEN, all_users=True, state_filter="done" ) with mock.patch( - "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" + "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: page = six.next(iterator.pages)