uranusjr commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r3975813046
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -554,29 +624,78 @@ def _should_run_output_processing(self) -> bool:
def execute(self, context):
self.log.info("Executing: %s", self.sql)
- hook = self.get_db_hook()
- if self.split_statements is not None:
- extra_kwargs = {"split_statements": self.split_statements}
+ if self.deferrable:
+ if self.enforce_read_only:
+ is_write, reason = scan_for_writes(self.sql)
+ if is_write:
+ raise ValueError(
+ f"enforce_read_only=True but the SQL appears to
contain a write: {reason}. "
+ "Set enforce_read_only=False if the query is
idempotent, or deferrable=False "
+ "to run it on the worker."
+ )
+ self.defer(
+ trigger=SQLExecuteQueryTrigger(
+ sql=self.sql,
+ conn_id=self.conn_id,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ fetch_results=self._should_run_output_processing() or
self.requires_result_fetch,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ read_only=self.enforce_read_only,
+ ),
+ method_name="execute_complete",
+ )
else:
- extra_kwargs = {}
- output = hook.run(
- sql=self.sql,
- autocommit=self.autocommit,
- parameters=self.parameters,
- handler=self.handler
- if self._should_run_output_processing() or
self.requires_result_fetch
- else None,
- return_last=self.return_last,
- **extra_kwargs,
- )
- if not self._should_run_output_processing():
+ hook = self.get_db_hook()
+ if self.split_statements is not None:
+ extra_kwargs = {"split_statements": self.split_statements}
+ else:
+ extra_kwargs = {}
+ output = hook.run(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=self.handler
+ if self._should_run_output_processing() or
self.requires_result_fetch
+ else None,
+ return_last=self.return_last,
+ **extra_kwargs,
+ )
+ if not self._should_run_output_processing():
+ return None
+ if return_single_query_results(self.sql, self.return_last,
self.split_statements):
+ # For simplicity, we pass always list as input to
_process_output, regardless if
+ # single query results are going to be returned, and we return
the first element
+ # of the list in this case from the (always) list returned by
_process_output
+ return self._process_output([output], hook.descriptions)[-1]
+ result = self._process_output(output, hook.descriptions)
+ self.log.info("result: %s", result)
+ return result
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> Any:
+ if event is None:
+ raise RuntimeError("Unknown error in SQLExecuteQueryTrigger")
+ if event.get("status") == "error":
+ raise RuntimeError(event.get("message", "Unknown error in
SQLExecuteQueryTrigger"))
+ self.log.info("SQL query executed successfully.")
+ results = event.get("results")
+ if not self._should_run_output_processing() or results is None:
return None
+ descriptions: list[Sequence[Sequence] | None] =
event.get("descriptions") or []
if return_single_query_results(self.sql, self.return_last,
self.split_statements):
- # For simplicity, we pass always list as input to _process_output,
regardless if
- # single query results are going to be returned, and we return the
first element
- # of the list in this case from the (always) list returned by
_process_output
- return self._process_output([output], hook.descriptions)[-1]
- result = self._process_output(output, hook.descriptions)
+ # The trigger returns the rows of a single statement; apply the
handler on the worker, where
+ # it is defined, over a replayed cursor rather than in the
triggerer's event loop.
+ rows = [results] if isinstance(results, str) else results
Review Comment:
When would a fetched result be a bare str here?
--
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]
