Visitar URL original
feat: accept job object as argument to `get_job` and `cancel_job` (#617) · googleapis/google-cloud-python@053a498 · GitHub
Skip to content

Commit 053a498

Browse files
tswasttswast
authored
feat: accept job object as argument to get_job and cancel_job (#617)
This allows one to more easily cancel or get updated metadata for an existing job from the client class. Ensures that project ID and location are correctly populated. Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #616 🦕
1 parent 1762f29 commit 053a498

3 files changed

Lines changed: 93 additions & 24 deletions

File tree

‎packages/google-cloud-bigquery/google/cloud/bigquery/client.py‎

Lines changed: 59 additions & 4 deletions
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

‎packages/google-cloud-bigquery/tests/system/test_client.py‎

Lines changed: 6 additions & 5 deletions
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

‎packages/google-cloud-bigquery/tests/unit/test_client.py‎

Lines changed: 28 additions & 15 deletions
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)