potiuk commented on code in PR #76106:
URL: https://github.com/apache/airflow/pull/76106#discussion_r4435622987
##########
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:
**task state written through the public API is not Dag-author input**
The job name read back from `task_state_store` becomes `self.job_name` here
(and in `get_job_result`), and is then formatted unescaped into `DESCRIBE
SERVICE {...}`, `DROP SERVICE IF EXISTS {...}` and
`SYSTEM$GET_SERVICE_LOGS('{...}', ...)`. Task state can also be written through
the public REST API (`PUT .../taskStateStore`) with task-instance edit
permission, which is less than Dag authoring, so a stored value is not
necessarily something the operator wrote. Could the stored name be validated
against Snowflake identifier syntax before it reaches SQL (the names are
Snowflake-generated, so a strict check is safe), the `GET_SERVICE_LOGS`
argument be bound instead of formatted, and the DONE marker be tied to the job
name (e.g. `{job_name: "DONE"}`) so a marker set from outside cannot make the
operator skip the job either?
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -123,16 +139,22 @@ def __init__(
poll_interval: int = 10,
snowflake_conn_id: str = "snowflake_default",
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ durable: bool | None = None,
+ cancel_on_kill: bool = True,
Review Comment:
**cancel_on_kill default vs durable default**
Thanks for adding `cancel_on_kill`. With both `durable` and `cancel_on_kill`
defaulting to `True`, a graceful SIGTERM (pod eviction, worker shutdown) still
drops a durable job and the retry starts fresh. Defaulting `cancel_on_kill` to
`not durable`, or documenting that durable jobs need `cancel_on_kill=False` to
survive preemption, would make the default match the durable 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]
