justinpakzad commented on code in PR #76106:
URL: https://github.com/apache/airflow/pull/76106#discussion_r4433637357
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -207,14 +254,25 @@ def _poll_for_status(self) -> str:
if self.drop_on_completion:
self._drop_service()
raise TimeoutError(f"Job {self.job_name} did not reach a
terminal status before the timeout.")
- response = self._run_one(f"DESCRIBE SERVICE {self.job_name}",
return_dictionaries=True)
- status = response.get("status")
+ status = self._describe_status(self.job_name)
if status in TERMINAL_STATUSES:
- return status
+ # get_job_result is skipped when the mixin reconnects to a
still-running job, so
+ # finalize here.
+ self._handle_final_status(status=status)
Review Comment:
The DONE marker is now written to the task state store right before
`_handle_final_status` and `get_job_status` was updated to check the store.
Tests were added for these changes as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
