justinpakzad commented on code in PR #76106:
URL: https://github.com/apache/airflow/pull/76106#discussion_r4433636576
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -191,14 +211,41 @@ def _run_one(self, sql: str, return_dictionaries: bool =
False) -> Any:
"""Run a single statement that returns one row via
fetch_one_handler."""
return self._hook.run(sql, handler=fetch_one_handler,
return_dictionaries=return_dictionaries)
- def _submit_job(self) -> str:
+ def _describe_status(self, external_id: JsonValue) -> str:
+ """Describe the job's current status."""
+ response = self._run_one(f"DESCRIBE SERVICE {external_id}",
return_dictionaries=True)
+ return response.get("status")
+
+ def submit_job(self, context: Context) -> str:
"""Submit the job and return the name."""
response = self._run_one(self._build_sql())
- job_name = response[0].split("'")[1]
- return job_name
+ self.job_name = response[0].split("'")[1]
+ if not self.job_name:
+ raise RuntimeError("Job name was not returned")
+ return self.job_name
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ """Return NOT_FOUND when the service no longer exists, otherwise the
current status."""
+ try:
+ return self._describe_status(external_id)
+ except ProgrammingError as e:
+ if e.errno == OBJECT_NOT_EXIST_ERROR_CODE:
+ return NOT_FOUND_STATUS
+ raise
+
+ def is_job_active(self, status: str) -> bool:
Review Comment:
Yea, good call. `is_job_active` now treats only PENDING/RUNNING as
reconnectable.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -191,14 +211,41 @@ def _run_one(self, sql: str, return_dictionaries: bool =
False) -> Any:
"""Run a single statement that returns one row via
fetch_one_handler."""
return self._hook.run(sql, handler=fetch_one_handler,
return_dictionaries=return_dictionaries)
- def _submit_job(self) -> str:
+ def _describe_status(self, external_id: JsonValue) -> str:
+ """Describe the job's current status."""
+ response = self._run_one(f"DESCRIBE SERVICE {external_id}",
return_dictionaries=True)
+ return response.get("status")
+
+ def submit_job(self, context: Context) -> str:
"""Submit the job and return the name."""
response = self._run_one(self._build_sql())
- job_name = response[0].split("'")[1]
- return job_name
+ self.job_name = response[0].split("'")[1]
+ if not self.job_name:
+ raise RuntimeError("Job name was not returned")
+ return self.job_name
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ """Return NOT_FOUND when the service no longer exists, otherwise the
current status."""
+ try:
+ return self._describe_status(external_id)
+ except ProgrammingError as e:
+ if e.errno == OBJECT_NOT_EXIST_ERROR_CODE:
+ return NOT_FOUND_STATUS
+ raise
+
+ def is_job_active(self, status: str) -> bool:
+ """Return True while the job is still running."""
+ return status in NON_TERMINAL_STATUSES
+
+ def is_job_succeeded(self, status: str) -> bool:
+ """Return True when the job has completed successfully."""
+ return status == SnowparkContainerJobStatus.DONE
+
+ def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
+ """Poll the job until it reaches a terminal state and handle the final
status."""
+ # On reconnect the mixin skips submit_job, so set the job name from
the external id here.
+ self.job_name = cast("str", external_id)
- def _poll_for_status(self) -> str:
- """Poll until the job reaches a terminal state."""
status = None
end_time = time.monotonic() + self.timeout
Review Comment:
Noted in the docs.
--
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]
