amoghrajesh commented on code in PR #69477:
URL: https://github.com/apache/airflow/pull/69477#discussion_r3569204630
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -507,6 +528,73 @@ 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)
+ # On reconnect, execute_query (the only thing that normally populates
this) never ran
+ # on this hook instance -- sync it so OpenLineage's
get_openlineage_database_specific_lineage,
+ # which reads hook.query_ids (not the operator's), doesn't silently
produce no lineage.
+ self._hook.query_ids = self.query_ids
+ # Push before polling, not after, so the handles are recorded even if
a statement
+ # errors below.
+ if self.do_xcom_push and context is not None:
+ context["ti"].xcom_push(key="query_ids", value=self.query_ids)
+ while True:
+ statement_status = self.poll_on_queries()
+ if statement_status["error"]:
+ raise RuntimeError(str(statement_status["error"]))
+ if not statement_status["running"]:
+ break
+ # On reconnect, the mixin calls poll_until_complete alone --
get_job_result is never
+ # invoked in that case -- so the output must be fetched here too, not
left to
+ # get_job_result. Fresh submit calls both; the flag stops
get_job_result from
+ # fetching (and pushing xcoms) a second time.
+ self._hook.check_query_output(self.query_ids)
+ self._poll_until_complete_ran = True
+
+ def get_job_result(self, external_id: JsonValue, context: Context) -> None:
+ self.query_ids = cast("list[str]", external_id)
+ # Same reconnect-hook gap as poll_until_complete -- see the comment
there. This path
+ # hits it too, since the already-succeeded case never calls
poll_until_complete either.
+ self._hook.query_ids = self.query_ids
+ if getattr(self, "_poll_until_complete_ran", False):
+ return
+ # The already-succeeded retry path skips submit_job and
poll_until_complete entirely,
+ # so push the query_ids xcom and fetch output here for parity with the
normal path.
+ if self.do_xcom_push and context is not None:
+ context["ti"].xcom_push(key="query_ids", value=self.query_ids)
+ self._hook.check_query_output(self.query_ids)
Review Comment:
Good observation, handled as suggested: d04f4e0f171
--
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]
