Skip to content

Commit 81b85eb

Browse files
pankajastropankajastro
authored
Fix BigQueryInsertJobOperator not exiting deferred state (#31591)
1 parent fcbbf47 commit 81b85eb

4 files changed

Lines changed: 59 additions & 88 deletions

File tree

‎airflow/providers/google/cloud/hooks/bigquery.py‎

Lines changed: 9 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3103,29 +3103,16 @@ async def get_job_instance(
31033103
with await self.service_file_as_context() as f:
31043104
return Job(job_id=job_id, project=project_id, service_file=f, session=cast(Session, session))
31053105

3106-
async def get_job_status(
3107-
self,
3108-
job_id: str | None,
3109-
project_id: str | None = None,
3110-
) -> str | None:
3111-
"""Poll for job status asynchronously using gcloud-aio.
3112-
3113-
Note that an OSError is raised when Job results are still pending.
3114-
Exception means that Job finished with errors
3115-
"""
3106+
async def get_job_status(self, job_id: str | None, project_id: str | None = None) -> str:
31163107
async with ClientSession() as s:
3117-
try:
3118-
self.log.info("Executing get_job_status...")
3119-
job_client = await self.get_job_instance(project_id, job_id, s)
3120-
job_status_response = await job_client.result(cast(Session, s))
3121-
if job_status_response:
3122-
job_status = "success"
3123-
except OSError:
3124-
job_status = "pending"
3125-
except Exception as e:
3126-
self.log.info("Query execution finished with errors...")
3127-
job_status = str(e)
3128-
return job_status
3108+
job_client = await self.get_job_instance(project_id, job_id, s)
3109+
job = await job_client.get_job()
3110+
status = job.get("status", {})
3111+
if status["state"] == "DONE":
3112+
if "errorResult" in status:
3113+
return "error"
3114+
return "success"
3115+
return status["state"].lower()
31293116

31303117
async def get_job_output(
31313118
self,

‎airflow/providers/google/cloud/triggers/bigquery.py‎

Lines changed: 31 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -72,31 +72,29 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
7272
)
7373

7474
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
75+
"""Gets current job execution status and yields a TriggerEvent."""
7576
"""Gets current job execution status and yields a TriggerEvent."""
7677
hook = self._get_async_hook()
7778
while True:
7879
try:
79-
# Poll for job execution status
80-
response_from_hook = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
81-
self.log.debug("Response from hook: %s", response_from_hook)
82-
83-
if response_from_hook == "success":
80+
job_status = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
81+
if job_status == "success":
8482
yield TriggerEvent(
8583
{
8684
"job_id": self.job_id,
87-
"status": "success",
85+
"status": job_status,
8886
"message": "Job completed",
8987
}
9088
)
9189
return
92-
elif response_from_hook == "pending":
93-
self.log.info("Query is still running...")
94-
self.log.info("Sleeping for %s seconds.", self.poll_interval)
95-
await asyncio.sleep(self.poll_interval)
96-
else:
97-
yield TriggerEvent({"status": "error", "message": response_from_hook})
90+
elif job_status == "error":
91+
yield TriggerEvent({"status": "error"})
9892
return
99-
93+
else:
94+
self.log.info(
95+
"Bigquery job status is %s. Sleeping for %s seconds.", job_status, self.poll_interval
96+
)
97+
await asyncio.sleep(self.poll_interval)
10098
except Exception as e:
10199
self.log.exception("Exception occurred while checking for query completion")
102100
yield TriggerEvent({"status": "error", "message": str(e)})
@@ -129,8 +127,8 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
129127
while True:
130128
try:
131129
# Poll for job execution status
132-
response_from_hook = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
133-
if response_from_hook == "success":
130+
job_status = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
131+
if job_status == "success":
134132
query_results = await hook.get_job_output(job_id=self.job_id, project_id=self.project_id)
135133

136134
records = hook.get_records(query_results)
@@ -154,14 +152,14 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
154152
}
155153
)
156154
return
157-
158-
elif response_from_hook == "pending":
159-
self.log.info("Query is still running...")
160-
self.log.info("Sleeping for %s seconds.", self.poll_interval)
161-
await asyncio.sleep(self.poll_interval)
162-
else:
163-
yield TriggerEvent({"status": "error", "message": response_from_hook})
155+
elif job_status == "error":
156+
yield TriggerEvent({"status": "error", "message": job_status})
164157
return
158+
else:
159+
self.log.info(
160+
"Bigquery job status is %s. Sleeping for %s seconds.", job_status, self.poll_interval
161+
)
162+
await asyncio.sleep(self.poll_interval)
165163
except Exception as e:
166164
self.log.exception("Exception occurred while checking for query completion")
167165
yield TriggerEvent({"status": "error", "message": str(e)})
@@ -201,26 +199,27 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
201199
while True:
202200
try:
203201
# Poll for job execution status
204-
response_from_hook = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
205-
if response_from_hook == "success":
202+
job_status = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
203+
if job_status == "success":
206204
query_results = await hook.get_job_output(job_id=self.job_id, project_id=self.project_id)
207205
records = hook.get_records(query_results=query_results, as_dict=self.as_dict)
208-
self.log.debug("Response from hook: %s", response_from_hook)
206+
self.log.debug("Response from hook: %s", job_status)
209207
yield TriggerEvent(
210208
{
211209
"status": "success",
212-
"message": response_from_hook,
210+
"message": job_status,
213211
"records": records,
214212
}
215213
)
216214
return
217-
elif response_from_hook == "pending":
218-
self.log.info("Query is still running...")
219-
self.log.info("Sleeping for %s seconds.", self.poll_interval)
220-
await asyncio.sleep(self.poll_interval)
221-
else:
222-
yield TriggerEvent({"status": "error", "message": response_from_hook})
215+
elif job_status == "error":
216+
yield TriggerEvent({"status": "error"})
223217
return
218+
else:
219+
self.log.info(
220+
"Bigquery job status is %s. Sleeping for %s seconds.", job_status, self.poll_interval
221+
)
222+
await asyncio.sleep(self.poll_interval)
224223
except Exception as e:
225224
self.log.exception("Exception occurred while checking for query completion")
226225
yield TriggerEvent({"status": "error", "message": str(e)})

‎tests/providers/google/cloud/hooks/test_bigquery.py‎

Lines changed: 11 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -2139,34 +2139,23 @@ async def test_get_job_instance(self, mock_session):
21392139
result = await hook.get_job_instance(project_id=PROJECT_ID, job_id=JOB_ID, session=mock_session)
21402140
assert isinstance(result, Job)
21412141

2142+
@pytest.mark.parametrize(
2143+
"job_status, expected",
2144+
[
2145+
({"status": {"state": "DONE"}}, "success"),
2146+
({"status": {"state": "DONE", "errorResult": "Timeout"}}, "error"),
2147+
({"status": {"state": "running"}}, "running"),
2148+
],
2149+
)
21422150
@pytest.mark.asyncio
21432151
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_instance")
2144-
async def test_get_job_status_success(self, mock_job_instance):
2152+
async def test_get_job_status(self, mock_job_instance, job_status, expected):
21452153
hook = BigQueryAsyncHook()
21462154
mock_job_client = AsyncMock(Job)
21472155
mock_job_instance.return_value = mock_job_client
2148-
response = "success"
2149-
mock_job_instance.return_value.result.return_value = response
2156+
mock_job_instance.return_value.get_job.return_value = job_status
21502157
resp = await hook.get_job_status(job_id=JOB_ID, project_id=PROJECT_ID)
2151-
assert resp == response
2152-
2153-
@pytest.mark.asyncio
2154-
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_instance")
2155-
async def test_get_job_status_oserror(self, mock_job_instance):
2156-
"""Assets that the BigQueryAsyncHook returns a pending response when OSError is raised"""
2157-
mock_job_instance.return_value.result.side_effect = OSError()
2158-
hook = BigQueryAsyncHook()
2159-
job_status = await hook.get_job_status(job_id=JOB_ID, project_id=PROJECT_ID)
2160-
assert job_status == "pending"
2161-
2162-
@pytest.mark.asyncio
2163-
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_instance")
2164-
async def test_get_job_status_exception(self, mock_job_instance, caplog):
2165-
"""Assets that the logging is done correctly when BigQueryAsyncHook raises Exception"""
2166-
mock_job_instance.return_value.result.side_effect = Exception()
2167-
hook = BigQueryAsyncHook()
2168-
await hook.get_job_status(job_id=JOB_ID, project_id=PROJECT_ID)
2169-
assert "Query execution finished with errors..." in caplog.text
2158+
assert resp == expected
21702159

21712160
@pytest.mark.asyncio
21722161
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_instance")

‎tests/providers/google/cloud/triggers/test_bigquery.py‎

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ async def test_bigquery_insert_job_trigger_running(self, mock_job_instance, capl
180180

181181
mock_job_client = AsyncMock(Job)
182182
mock_job_instance.return_value = mock_job_client
183-
mock_job_instance.return_value.result.side_effect = OSError
183+
mock_job_instance.return_value.get_job.return_value = {"status": {"state": "running"}}
184184
caplog.set_level(logging.INFO)
185185

186186
task = asyncio.create_task(insert_job_trigger.run().__anext__())
@@ -189,8 +189,7 @@ async def test_bigquery_insert_job_trigger_running(self, mock_job_instance, capl
189189
# TriggerEvent was not returned
190190
assert task.done() is False
191191

192-
assert "Query is still running..." in caplog.text
193-
assert f"Sleeping for {POLLING_PERIOD_SECONDS} seconds." in caplog.text
192+
assert "Bigquery job status is running. Sleeping for 4.0 seconds." in caplog.text
194193

195194
# Prevents error when task is destroyed while in "pending" state
196195
asyncio.get_event_loop().stop()
@@ -205,7 +204,7 @@ async def test_bigquery_op_trigger_terminated(self, mock_job_status, caplog, ins
205204

206205
generator = insert_job_trigger.run()
207206
actual = await generator.asend(None)
208-
assert TriggerEvent({"status": "error", "message": "error"}) == actual
207+
assert TriggerEvent({"status": "error"}) == actual
209208

210209
@pytest.mark.asyncio
211210
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
@@ -241,7 +240,7 @@ async def test_bigquery_get_data_trigger_running(self, mock_job_instance, caplog
241240

242241
mock_job_client = AsyncMock(Job)
243242
mock_job_instance.return_value = mock_job_client
244-
mock_job_instance.return_value.result.side_effect = OSError
243+
mock_job_instance.return_value.get_job.return_value = {"status": {"state": "RUNNING"}}
245244
caplog.set_level(logging.INFO)
246245

247246
task = asyncio.create_task(get_data_trigger.run().__anext__())
@@ -250,8 +249,7 @@ async def test_bigquery_get_data_trigger_running(self, mock_job_instance, caplog
250249
# TriggerEvent was not returned
251250
assert task.done() is False
252251

253-
assert "Query is still running..." in caplog.text
254-
assert f"Sleeping for {POLLING_PERIOD_SECONDS} seconds." in caplog.text
252+
assert "Bigquery job status is running. Sleeping for 4.0 seconds." in caplog.text
255253

256254
# Prevents error when task is destroyed while in "pending" state
257255
asyncio.get_event_loop().stop()
@@ -266,7 +264,7 @@ async def test_bigquery_get_data_trigger_terminated(self, mock_job_status, caplo
266264

267265
generator = get_data_trigger.run()
268266
actual = await generator.asend(None)
269-
assert TriggerEvent({"status": "error", "message": "error"}) == actual
267+
assert TriggerEvent({"status": "error"}) == actual
270268

271269
@pytest.mark.asyncio
272270
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
@@ -336,17 +334,15 @@ async def test_bigquery_check_trigger_running(self, mock_job_instance, caplog, c
336334

337335
mock_job_client = AsyncMock(Job)
338336
mock_job_instance.return_value = mock_job_client
339-
mock_job_instance.return_value.result.side_effect = OSError
340-
caplog.set_level(logging.INFO)
337+
mock_job_instance.return_value.get_job.return_value = {"status": {"state": "running"}}
341338

342339
task = asyncio.create_task(check_trigger.run().__anext__())
343340
await asyncio.sleep(0.5)
344341

345342
# TriggerEvent was not returned
346343
assert task.done() is False
347344

348-
assert "Query is still running..." in caplog.text
349-
assert f"Sleeping for {POLLING_PERIOD_SECONDS} seconds." in caplog.text
345+
assert "Bigquery job status is running. Sleeping for 4.0 seconds." in caplog.text
350346

351347
# Prevents error when task is destroyed while in "pending" state
352348
asyncio.get_event_loop().stop()

0 commit comments

Comments
 (0)