RNHTTR commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r3501298787
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -552,31 +616,82 @@ def _process_output(
def _should_run_output_processing(self) -> bool:
return self.do_xcom_push
+ def _get_handler_import_path(self) -> str:
Review Comment:
Does this get called anywhere?
##########
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:
What happens if the triggerer is killed in the middle of a run? Will this
trigger be re-executed? Is there a chance that multiple instances of the same
query are invoked?
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -409,7 +409,7 @@ def __init__(
"session_parameters": session_parameters,
**hook_params,
}
- super().__init__(conn_id=snowflake_conn_id, **kwargs) # pragma: no
cover
+ super().__init__(conn_id=snowflake_conn_id, deferrable=deferrable,
**kwargs) # pragma: no cover
Review Comment:
is this intentional?
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -17,21 +17,32 @@
# under the License.
from __future__ import annotations
+import importlib
+import sys
from typing import TYPE_CHECKING
+from asgiref.sync import sync_to_async
+
from airflow.providers.common.compat.sdk import AirflowException, BaseHook
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS
+from airflow.providers.common.sql.hooks.handlers import fetch_all_handler
from airflow.providers.common.sql.hooks.sql import DbApiHook
from airflow.triggers.base import BaseTrigger, TriggerEvent
if TYPE_CHECKING:
from collections.abc import AsyncIterator
from typing import Any
+from collections.abc import (
+ Iterable,
+ Mapping,
+ Sequence,
+)
-class SQLExecuteQueryTrigger(BaseTrigger):
+
+class SQLGenericTransferTrigger(BaseTrigger):
Review Comment:
This is a breaking change, right? If so, we'll need to bump the major
version of this provider.
--
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]
