justinpakzad commented on code in PR #76106:
URL: https://github.com/apache/airflow/pull/76106#discussion_r4436277415
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -191,14 +217,44 @@ 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 DONE if a prior try stored it, NOT_FOUND if the service is
gone, else the current status."""
+ if (task_state_store := context.get("task_state_store")) is not None:
+ if task_state_store.get(self._job_status_key) ==
SnowparkContainerJobStatus.DONE:
+ return SnowparkContainerJobStatus.DONE
+ 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 (SnowparkContainerJobStatus.PENDING,
SnowparkContainerJobStatus.RUNNING)
+
+ 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)
Review Comment:
Very good point. The job name now goes through a `_validate_job_name` helper
which validates against Snowflake identifier syntax, the `GET_SERVICE_LOGS`
args are now bounded, and the marker is tied to a job name.
--
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]
