diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py index a4e57ef3e..7a4478551 100644 --- a/google/cloud/bigquery/job.py +++ b/google/cloud/bigquery/job.py @@ -284,19 +284,26 @@ class _AsyncJob(google.api_core.future.polling.PollingFuture): """ def __init__(self, job_id, client): super(_AsyncJob, self).__init__() + + # The job reference can be either a plain job ID or the full resource. + # Populate the properties dictionary consistently depending on what has + # been passed in. job_ref = job_id if not isinstance(job_id, _JobReference): job_ref = _JobReference(job_id, client.project, None) - self._job_ref = job_ref + self._properties = { + 'jobReference': job_ref._to_api_repr(), + } + self._client = client - self._properties = {} self._result_set = False self._completion_lock = threading.Lock() @property def job_id(self): """str: ID of the job.""" - return self._job_ref.job_id + return _helpers._get_sub_prop( + self._properties, ['jobReference', 'jobId']) @property def project(self): @@ -305,12 +312,14 @@ def project(self): :rtype: str :returns: the project (derived from the client). """ - return self._job_ref.project + return _helpers._get_sub_prop( + self._properties, ['jobReference', 'projectId']) @property def location(self): """str: Location where the job runs.""" - return self._job_ref.location + return _helpers._get_sub_prop( + self._properties, ['jobReference', 'location']) def _require_client(self, client): """Check client or verify over-ride. @@ -481,7 +490,7 @@ def _set_properties(self, api_response): self._properties.clear() self._properties.update(cleaned) - self._copy_configuration_properties(cleaned['configuration']) + self._copy_configuration_properties(cleaned.get('configuration', {})) # For Future interface self._set_future_result() @@ -1337,7 +1346,7 @@ def _build_resource(self): self.destination.to_api_repr()) return { - 'jobReference': self._job_ref._to_api_repr(), + 'jobReference': self._properties['jobReference'], 'configuration': configuration, } @@ -1523,7 +1532,7 @@ def _build_resource(self): }) return { - 'jobReference': self._job_ref._to_api_repr(), + 'jobReference': self._properties['jobReference'], 'configuration': configuration, } @@ -1737,7 +1746,7 @@ def _build_resource(self): self.destination_uris) return { - 'jobReference': self._job_ref._to_api_repr(), + 'jobReference': self._properties['jobReference'], 'configuration': configuration, } @@ -2330,7 +2339,7 @@ def _build_resource(self): configuration = self._configuration.to_api_repr() resource = { - 'jobReference': self._job_ref._to_api_repr(), + 'jobReference': self._properties['jobReference'], 'configuration': configuration, } configuration['query']['query'] = self.query @@ -3020,8 +3029,12 @@ def from_api_repr(cls, resource, client): Returns: UnknownJob: Job corresponding to the resource. """ - job_ref = _JobReference._from_api_repr( - resource.get('jobReference', {'projectId': client.project})) + job_ref_properties = resource.get( + 'jobReference', {'projectId': client.project}) + job_ref = _JobReference._from_api_repr(job_ref_properties) job = cls(job_ref, client) + # Populate the job reference with the project, even if it has been + # redacted, because we know it should equal that of the request. + resource['jobReference'] = job_ref_properties job._properties = resource return job diff --git a/tests/system.py b/tests/system.py index 6662f584c..a94a67219 100644 --- a/tests/system.py +++ b/tests/system.py @@ -666,7 +666,7 @@ def test_load_table_from_file_w_explicit_location(self): client.get_job(job_id, location='US') load_job_us = client.get_job(job_id) - load_job_us._job_ref._properties['location'] = 'US' + load_job_us._properties['jobReference']['location'] = 'US' self.assertFalse(load_job_us.exists()) with self.assertRaises(NotFound): load_job_us.reload() diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 85f882eee..ee7bb7ca8 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -2740,6 +2740,39 @@ def test_query_w_client_location(self): data=resource, ) + def test_query_detect_location(self): + query = 'select count(*) from persons' + resource_location = 'EU' + resource = { + 'jobReference': { + 'projectId': self.PROJECT, + # Location not set in request, but present in the response. + 'location': resource_location, + 'jobId': 'some-random-id', + }, + 'configuration': { + 'query': { + 'query': query, + 'useLegacySql': False, + }, + }, + } + creds = _make_credentials() + http = object() + client = self._make_one(project=self.PROJECT, credentials=creds, + _http=http) + conn = client._connection = _make_connection(resource) + + job = client.query(query) + + self.assertEqual(job.location, resource_location) + + # Check that request did not contain a location. + conn.api_request.assert_called_once() + _, req = conn.api_request.call_args + sent = req['data'] + self.assertIsNone(sent['jobReference'].get('location')) + def test_query_w_udf_resources(self): from google.cloud.bigquery.job import QueryJob from google.cloud.bigquery.job import QueryJobConfig diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 0df830a2d..f3a57439c 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -156,7 +156,15 @@ def test_ctor_w_bare_job_id(self): self.assertEqual(job.project, self.PROJECT) self.assertIsNone(job.location) self.assertIs(job._client, client) - self.assertEqual(job._properties, {}) + self.assertEqual( + job._properties, + { + 'jobReference': { + 'projectId': self.PROJECT, + 'jobId': self.JOB_ID, + }, + } + ) self.assertIsInstance(job._completion_lock, type(threading.Lock())) self.assertEqual( job.path, @@ -174,7 +182,16 @@ def test_ctor_w_job_ref(self): self.assertEqual(job.project, self.PROJECT) self.assertEqual(job.location, self.LOCATION) self.assertIs(job._client, client) - self.assertEqual(job._properties, {}) + self.assertEqual( + job._properties, + { + 'jobReference': { + 'projectId': self.PROJECT, + 'location': self.LOCATION, + 'jobId': self.JOB_ID, + }, + } + ) self.assertFalse(job._result_set) self.assertIsInstance(job._completion_lock, type(threading.Lock())) self.assertEqual( @@ -359,6 +376,7 @@ def _set_properties_job(self): job._copy_configuration_properties = mock.Mock() job._set_future_result = mock.Mock() job._properties = { + 'jobReference': job._properties['jobReference'], 'foo': 'bar', } return job @@ -577,7 +595,7 @@ def test_exists_defaults_miss(self): from google.cloud.bigquery.retry import DEFAULT_RETRY job = self._set_properties_job() - job._job_ref._properties['location'] = self.LOCATION + job._properties['jobReference']['location'] = self.LOCATION call_api = job._client._call_api = mock.Mock() call_api.side_effect = NotFound('testing') @@ -636,7 +654,7 @@ def test_reload_defaults(self): } } job = self._set_properties_job() - job._job_ref._properties['location'] = self.LOCATION + job._properties['jobReference']['location'] = self.LOCATION call_api = job._client._call_api = mock.Mock() call_api.return_value = resource @@ -693,7 +711,7 @@ def test_cancel_defaults(self): } response = {'job': resource} job = self._set_properties_job() - job._job_ref._properties['location'] = self.LOCATION + job._properties['jobReference']['location'] = self.LOCATION connection = job._client._connection = _make_connection(response) self.assertTrue(job.cancel())