karenbraganz commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r3580568882
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -103,3 +114,142 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
except Exception as e:
self.log.exception("An error occurred: %s", e)
yield TriggerEvent({"status": "failure", "message": str(e)})
+
+
+class SQLExecuteQueryTrigger(BaseTrigger):
+ """
+ A SQL trigger that executes SQL code in async mode.
+
+ The query runs in the triggerer, but no user code does: when
``fetch_results`` is set the rows are
+ fetched with the built-in :func:`fetch_all_handler` and returned, together
with the cursor
+ descriptions, in the ``TriggerEvent``. Any user-provided ``handler`` is
applied on the worker in
+ ``SQLExecuteQueryOperator.execute_complete`` -- keeping user code out of
the triggerer's event loop
+ and out of its (bundle-less) import path.
+
+ :param sql: the sql statement to be executed (str) or a list of sql
statements to execute
+ :param conn_id: the connection ID used to connect to the database
+ :param fetch_results: whether the query results should be fetched and
returned to the worker
+ """
+
+ def __init__(
+ self,
+ sql: str | Iterable[str],
+ conn_id: str,
+ autocommit: bool,
+ split_statements: bool,
+ return_last: bool,
+ parameters: Iterable[Any] | Mapping[str, Any] | None = None,
+ fetch_results: bool = False,
+ ):
+ super().__init__()
+ self.sql = sql
+ self.conn_id = conn_id
+ self.autocommit = autocommit
+ self.parameters = parameters
+ self.fetch_results = fetch_results
+ self.split_statements = split_statements
+ self.return_last = return_last
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize the SQLExecuteQueryTrigger arguments and classpath."""
+ return (
+ f"{self.__class__.__module__}.{self.__class__.__name__}",
+ {
+ "sql": self.sql,
+ "conn_id": self.conn_id,
+ "autocommit": self.autocommit,
+ "parameters": self.parameters,
+ "fetch_results": self.fetch_results,
+ "split_statements": self.split_statements,
+ "return_last": self.return_last,
+ },
+ )
+
+ @staticmethod
+ def _jsonsafe_descriptions(
+ descriptions: list[Sequence[Sequence] | None],
+ ) -> list[list[list[Any]] | None]:
+ """
+ Normalise cursor descriptions into a JSON-serializable form for the
``TriggerEvent``.
+
+ ``cursor.description`` is a sequence of column 7-tuples whose
``type_code`` can be a
+ driver-specific object that is not JSON-serializable; such values are
stringified while
+ JSON-native fields (column names, sizes, precision, ...) are preserved.
+ """
+ safe: list[list[list[Any]] | None] = []
+ for description in descriptions:
+ if description is None:
+ safe.append(None)
+ continue
+ safe.append(
+ [
+ [
+ field if isinstance(field, (str, int, float, bool,
type(None))) else str(field)
+ for field in column
+ ]
+ for column in description
+ ]
+ )
+ return safe
+
+ async def get_hook(self) -> DbApiHook:
+ """
+ Return DbApiHook.
+
+ :return: DbApiHook for this connection
+ """
+ connection = await sync_to_async(BaseHook.get_connection)(self.conn_id)
+ hook = await sync_to_async(connection.get_hook)()
+ if not isinstance(hook, DbApiHook) or not hasattr(hook, "run_async"):
+ raise AirflowException(
+ f"You are trying to use the SqlExecuteQueryOperator in
deferrable mode with {hook.__class__.__name__},"
+ f" but its provider does not support this. Please set
deferrable=False"
+ f" Got {hook.__class__.__name__} with class hierarchy:
{hook.__class__.mro()}"
+ )
+ return hook
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ try:
+ hook = await self.get_hook()
+
+ self.log.info("Extracting data from %s", self.conn_id)
+ self.log.info("Executing: \n %s", self.sql)
+
+ if self.fetch_results:
+ # Fetch the raw rows with the built-in handler and return them
with the cursor
+ # descriptions; the operator applies any user handler on the
worker.
+ results = await hook.run_async(
Review Comment:
I looked into this. There isn't a straightforward way to implement this for
Postgres. It requires the client side connection to remain alive in order to
continue running the query. So we cannot open a connection from the worker to
start the query, then close that, and re-open a new connection on the
triggerer. Your second suggestion seems doable though:
> We could also potentially block non-read-only queries by default and
include a parameter to do things like INSERT, UPDATE, DELETE, etc with a
docstring explaining the risks?
I can look into implementing it like that.
--
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]
