From a3686161a0b8a7876fb4551d66e9fbce8c2cb398 Mon Sep 17 00:00:00 2001 From: Lars Klingen Date: Tue, 7 Sep 2021 18:30:29 -0300 Subject: [PATCH] Pass bigquery job object to get_job Signed-off-by: Lars Klingen --- sdk/python/feast/infra/offline_stores/bigquery.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py index 91feb1312f..7ac69ce685 100644 --- a/sdk/python/feast/infra/offline_stores/bigquery.py +++ b/sdk/python/feast/infra/offline_stores/bigquery.py @@ -261,11 +261,10 @@ def block_until_done( if is_test: retry_cadence = 0.1 - def _wait_until_done(job_id): - if client.get_job(job_id).state in ["PENDING", "RUNNING"]: - raise BigQueryJobStillRunning(job_id=job_id) + def _wait_until_done(bq_job): + if client.get_job(bq_job).state in ["PENDING", "RUNNING"]: + raise BigQueryJobStillRunning(job_id=bq_job.job_id) - job_id = bq_job.job_id try: retryer = Retrying( wait=wait_fixed(retry_cadence), @@ -273,12 +272,12 @@ def _wait_until_done(job_id): retry=retry_if_exception_type(BigQueryJobStillRunning), reraise=True, ) - retryer(_wait_until_done, job_id) + retryer(_wait_until_done, bq_job) finally: - if client.get_job(job_id).state in ["PENDING", "RUNNING"]: - client.cancel_job(job_id) - raise BigQueryJobCancelled(job_id=job_id) + if client.get_job(bq_job).state in ["PENDING", "RUNNING"]: + client.cancel_job(bq_job) + raise BigQueryJobCancelled(job_id=bq_job.job_id) if bq_job.exception(): raise bq_job.exception()