dabla commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r3507389565
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -103,3 +112,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):
Review Comment:
I would prefer to see the GenericTransfer also use this
SQLExecuteQueryTrigger, which is better suited due to it's more async native
nature. I don't like the fact that we have now a dedicated trigger for the
GenericTransfer, I think it should be possible to reuse the new
SQLExecuteQueryTrigger.
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi:
##########
@@ -185,6 +186,27 @@ class DbApiHook(BaseHook):
@staticmethod
def get_openlineage_authority_part(connection, default_port: int | None =
None) -> str: ...
def get_db_log_messages(self, conn) -> None: ...
+ async def async_get_conn(self) -> Any: ...
+ @overload
+ async def run_async(
Review Comment:
I would tend more to call the method `arun `instead of `run_async`, just to
follow the same naming convention applied like in TaskSDK (`aget_connection`,
`aget_hook`, ...)
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -103,3 +112,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)
Review Comment:
You should use the common compat `get_async_connection` here.
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -103,3 +112,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)()
Review Comment:
I could maybe also add a `get_async_hook` helper method in `common.compat`,
so that you can use it here as well just like the `get_async_connection`, that
would be much cleaner as newer Airflow will now also support the native async
`BaseHook.aget_hook` method which was added in this
[PR](https://github.com/apache/airflow/pull/68506).
--
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]
