kaxil commented on code in PR #69477:
URL: https://github.com/apache/airflow/pull/69477#discussion_r3537522750
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -507,6 +528,64 @@ def poll_on_queries(self):
"running": statement_running_status,
}
+ def submit_job(self, context: Context) -> JsonValue:
+ """Submit the SQL for execution and return the resulting statement
handles."""
+ self.log.info("Executing: %s", self.sql)
+ self.query_ids = self._hook.execute_query(
+ self.sql, statement_count=self.statement_count,
bindings=self.bindings, timeout=self.timeout
+ )
+ self.log.info("List of query ids %s", self.query_ids)
+ return cast("JsonValue", self.query_ids)
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ """Aggregate the status of every handle into a single verdict for the
mixin."""
+ statuses = []
+ for query_id in cast("list[str]", external_id):
+ try:
+
statuses.append(self._hook.get_sql_api_query_status(query_id)["status"])
+ except requests.exceptions.HTTPError as e:
+ if e.response is not None and e.response.status_code == 404:
+ return "not_found"
+ raise
+ if "error" in statuses:
+ return "error"
+ if "running" in statuses:
+ return "running"
+ return "success"
+
+ def is_job_active(self, status: str) -> bool:
+ return status == "running"
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return status == "success"
+
+ def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
+ self.query_ids = cast("list[str]", external_id)
+ while True:
+ statement_status = self.poll_on_queries()
+ if statement_status["error"]:
+ raise RuntimeError(str(statement_status["error"]))
+ if not statement_status["running"]:
+ break
+ if self.do_xcom_push and context is not None:
+ context["ti"].xcom_push(key="query_ids", value=self.query_ids)
Review Comment:
On the failure path `query_ids` is never pushed to xcom. When a statement
errors, `poll_on_queries()` reports the error and this method raises
`RuntimeError` at line 567, before it reaches this `xcom_push`. That is the
case where the handles matter most, since you need them to look the failed
statements up in Snowflake.
The pre-PR non-deferrable path pushed `query_ids` immediately after
`execute_query` (before the poll loop), so the handles were recorded even when
a query failed. This moves the push after the loop, so a failure now loses them.
Consider pushing the xcom near the top of `poll_until_complete` (right after
`self.query_ids = ...`, before the poll loop) so it lands on both the fresh and
reconnect paths and survives an error. Non-blocking.
--
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]
