File tree

3 files changed

+93
-24
lines changed

3 files changed

+93
-24
lines changed
Original file line numberDiff line numberDiff line change
@@ -1734,12 +1734,20 @@ def get_job(
17341734
https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/get
17351735
17361736
Args:
1737-
job_id (str): Unique job identifier.
1737+
job_id (Union[ \
1738+
str, \
1739+
google.cloud.bigquery.job.LoadJob, \
1740+
google.cloud.bigquery.job.CopyJob, \
1741+
google.cloud.bigquery.job.ExtractJob, \
1742+
google.cloud.bigquery.job.QueryJob \
1743+
]): Job identifier.
17381744
17391745
Keyword Arguments:
17401746
project (Optional[str]):
17411747
ID of the project which owns the job (defaults to the client's project).
1742-
location (Optional[str]): Location where the job was run.
1748+
location (Optional[str]):
1749+
Location where the job was run. Ignored if ``job_id`` is a job
1750+
object.
17431751
retry (Optional[google.api_core.retry.Retry]):
17441752
How to retry the RPC.
17451753
timeout (Optional[float]):
@@ -1757,6 +1765,10 @@ def get_job(
17571765
"""
17581766
extra_params = {"projection": "full"}
17591767

1768+
project, location, job_id = _extract_job_reference(
1769+
job_id, project=project, location=location
1770+
)
1771+
17601772
if project is None:
17611773
project = self.project
17621774

@@ -1791,12 +1803,20 @@ def cancel_job(
17911803
https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/cancel
17921804
17931805
Args:
1794-
job_id (str): Unique job identifier.
1806+
job_id (Union[ \
1807+
str, \
1808+
google.cloud.bigquery.job.LoadJob, \
1809+
google.cloud.bigquery.job.CopyJob, \
1810+
google.cloud.bigquery.job.ExtractJob, \
1811+
google.cloud.bigquery.job.QueryJob \
1812+
]): Job identifier.
17951813
17961814
Keyword Arguments:
17971815
project (Optional[str]):
17981816
ID of the project which owns the job (defaults to the client's project).
1799-
location (Optional[str]): Location where the job was run.
1817+
location (Optional[str]):
1818+
Location where the job was run. Ignored if ``job_id`` is a job
1819+
object.
18001820
retry (Optional[google.api_core.retry.Retry]):
18011821
How to retry the RPC.
18021822
timeout (Optional[float]):
@@ -1814,6 +1834,10 @@ def cancel_job(
18141834
"""
18151835
extra_params = {"projection": "full"}
18161836

1837+
project, location, job_id = _extract_job_reference(
1838+
job_id, project=project, location=location
1839+
)
1840+
18171841
if project is None:
18181842
project = self.project
18191843

@@ -3518,6 +3542,37 @@ def _item_to_table(iterator, resource):
35183542
return TableListItem(resource)
35193543

35203544

3545+
def _extract_job_reference(job, project=None, location=None):
3546+
"""Extract fully-qualified job reference from a job-like object.
3547+
3548+
Args:
3549+
job_id (Union[ \
3550+
str, \
3551+
google.cloud.bigquery.job.LoadJob, \
3552+
google.cloud.bigquery.job.CopyJob, \
3553+
google.cloud.bigquery.job.ExtractJob, \
3554+
google.cloud.bigquery.job.QueryJob \
3555+
]): Job identifier.
3556+
project (Optional[str]):
3557+
Project where the job was run. Ignored if ``job_id`` is a job
3558+
object.
3559+
location (Optional[str]):
3560+
Location where the job was run. Ignored if ``job_id`` is a job
3561+
object.
3562+
3563+
Returns:
3564+
Tuple[str, str, str]: ``(project, location, job_id)``
3565+
"""
3566+
if hasattr(job, "job_id"):
3567+
project = job.project
3568+
job_id = job.job_id
3569+
location = job.location
3570+
else:
3571+
job_id = job
3572+
3573+
return (project, location, job_id)
3574+
3575+
35213576
def _make_job_id(job_id, prefix=None):
35223577
"""Construct an ID for a new job.
35233578
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,9 @@ def test_get_service_account_email(self):
189189
def _create_bucket(self, bucket_name, location=None):
190190
storage_client = storage.Client()
191191
bucket = storage_client.bucket(bucket_name)
192-
retry_storage_errors(bucket.create)(location=location)
192+
retry_storage_errors(storage_client.create_bucket)(
193+
bucket_name, location=location
194+
)
193195
self.to_delete.append(bucket)
194196

195197
return bucket
@@ -872,7 +874,7 @@ def test_load_table_from_file_w_explicit_location(self):
872874
job_id = load_job.job_id
873875

874876
# Can get the job from the EU.
875-
load_job = client.get_job(job_id, location="EU")
877+
load_job = client.get_job(load_job)
876878
self.assertEqual(job_id, load_job.job_id)
877879
self.assertEqual("EU", load_job.location)
878880
self.assertTrue(load_job.exists())
@@ -889,7 +891,7 @@ def test_load_table_from_file_w_explicit_location(self):
889891

890892
# Can cancel the job from the EU.
891893
self.assertTrue(load_job.cancel())
892-
load_job = client.cancel_job(job_id, location="EU")
894+
load_job = client.cancel_job(load_job)
893895
self.assertEqual(job_id, load_job.job_id)
894896
self.assertEqual("EU", load_job.location)
895897

@@ -1204,8 +1206,7 @@ def test_query_w_timeout(self):
12041206
# Even though the query takes >1 second, the call to getQueryResults
12051207
# should succeed.
12061208
self.assertFalse(query_job.done(timeout=1))
1207-
1208-
Config.CLIENT.cancel_job(query_job.job_id, location=query_job.location)
1209+
self.assertIsNotNone(Config.CLIENT.cancel_job(query_job))
12091210

12101211
def test_query_w_page_size(self):
12111212
page_size = 45
Original file line numberDiff line numberDiff line change
@@ -2933,31 +2933,30 @@ def test_get_job_miss_w_explict_project(self):
29332933
conn = client._connection = make_connection()
29342934

29352935
with self.assertRaises(NotFound):
2936-
client.get_job(JOB_ID, project=OTHER_PROJECT, location=self.LOCATION)
2936+
client.get_job(JOB_ID, project=OTHER_PROJECT)
29372937

29382938
conn.api_request.assert_called_once_with(
29392939
method="GET",
29402940
path="/projects/OTHER_PROJECT/jobs/NONESUCH",
2941-
query_params={"projection": "full", "location": self.LOCATION},
2941+
query_params={"projection": "full"},
29422942
timeout=None,
29432943
)
29442944

29452945
def test_get_job_miss_w_client_location(self):
29462946
from google.cloud.exceptions import NotFound
29472947

2948-
OTHER_PROJECT = "OTHER_PROJECT"
29492948
JOB_ID = "NONESUCH"
29502949
creds = _make_credentials()
2951-
client = self._make_one(self.PROJECT, creds, location=self.LOCATION)
2950+
client = self._make_one("client-proj", creds, location="client-loc")
29522951
conn = client._connection = make_connection()
29532952

29542953
with self.assertRaises(NotFound):
2955-
client.get_job(JOB_ID, project=OTHER_PROJECT)
2954+
client.get_job(JOB_ID)
29562955

29572956
conn.api_request.assert_called_once_with(
29582957
method="GET",
2959-
path="/projects/OTHER_PROJECT/jobs/NONESUCH",
2960-
query_params={"projection": "full", "location": self.LOCATION},
2958+
path="/projects/client-proj/jobs/NONESUCH",
2959+
query_params={"projection": "full", "location": "client-loc"},
29612960
timeout=None,
29622961
)
29632962

@@ -2971,7 +2970,11 @@ def test_get_job_hit_w_timeout(self):
29712970
QUERY = "SELECT * from test_dataset:test_table"
29722971
ASYNC_QUERY_DATA = {
29732972
"id": "{}:{}".format(self.PROJECT, JOB_ID),
2974-
"jobReference": {"projectId": self.PROJECT, "jobId": "query_job"},
2973+
"jobReference": {
2974+
"projectId": "resource-proj",
2975+
"jobId": "query_job",
2976+
"location": "us-east1",
2977+
},
29752978
"state": "DONE",
29762979
"configuration": {
29772980
"query": {
@@ -2989,18 +2992,21 @@ def test_get_job_hit_w_timeout(self):
29892992
creds = _make_credentials()
29902993
client = self._make_one(self.PROJECT, creds)
29912994
conn = client._connection = make_connection(ASYNC_QUERY_DATA)
2995+
job_from_resource = QueryJob.from_api_repr(ASYNC_QUERY_DATA, client)
29922996

2993-
job = client.get_job(JOB_ID, timeout=7.5)
2997+
job = client.get_job(job_from_resource, timeout=7.5)
29942998

29952999
self.assertIsInstance(job, QueryJob)
29963000
self.assertEqual(job.job_id, JOB_ID)
3001+
self.assertEqual(job.project, "resource-proj")
3002+
self.assertEqual(job.location, "us-east1")
29973003
self.assertEqual(job.create_disposition, CreateDisposition.CREATE_IF_NEEDED)
29983004
self.assertEqual(job.write_disposition, WriteDisposition.WRITE_TRUNCATE)
29993005

30003006
conn.api_request.assert_called_once_with(
30013007
method="GET",
3002-
path="/projects/PROJECT/jobs/query_job",
3003-
query_params={"projection": "full"},
3008+
path="/projects/resource-proj/jobs/query_job",
3009+
query_params={"projection": "full", "location": "us-east1"},
30043010
timeout=7.5,
30053011
)
30063012

@@ -3049,25 +3055,32 @@ def test_cancel_job_hit(self):
30493055
QUERY = "SELECT * from test_dataset:test_table"
30503056
QUERY_JOB_RESOURCE = {
30513057
"id": "{}:{}".format(self.PROJECT, JOB_ID),
3052-
"jobReference": {"projectId": self.PROJECT, "jobId": "query_job"},
3058+
"jobReference": {
3059+
"projectId": "job-based-proj",
3060+
"jobId": "query_job",
3061+
"location": "asia-northeast1",
3062+
},
30533063
"state": "RUNNING",
30543064
"configuration": {"query": {"query": QUERY}},
30553065
}
30563066
RESOURCE = {"job": QUERY_JOB_RESOURCE}
30573067
creds = _make_credentials()
30583068
client = self._make_one(self.PROJECT, creds)
30593069
conn = client._connection = make_connection(RESOURCE)
3070+
job_from_resource = QueryJob.from_api_repr(QUERY_JOB_RESOURCE, client)
30603071

3061-
job = client.cancel_job(JOB_ID)
3072+
job = client.cancel_job(job_from_resource)
30623073

30633074
self.assertIsInstance(job, QueryJob)
30643075
self.assertEqual(job.job_id, JOB_ID)
3076+
self.assertEqual(job.project, "job-based-proj")
3077+
self.assertEqual(job.location, "asia-northeast1")
30653078
self.assertEqual(job.query, QUERY)
30663079

30673080
conn.api_request.assert_called_once_with(
30683081
method="POST",
3069-
path="/projects/PROJECT/jobs/query_job/cancel",
3070-
query_params={"projection": "full"},
3082+
path="/projects/job-based-proj/jobs/query_job/cancel",
3083+
query_params={"projection": "full", "location": "asia-northeast1"},
30713084
timeout=None,
30723085
)
30733086

0 commit comments

Comments
 (0)