amoghrajesh commented on code in PR #69477:
URL: https://github.com/apache/airflow/pull/69477#discussion_r3541131300
##########
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:
Good catch, handled it by pushing at top of `poll_until_complete` now:
544ffeec72
--
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]
