uranusjr commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r3975773925
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -57,49 +84,160 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
{
"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,
+ "read_only": self.read_only,
"hook_params": self.hook_params,
},
)
- def get_hook(self) -> DbApiHook:
+ @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 aget_hook(self) -> DbApiHook:
"""
Return DbApiHook.
:return: DbApiHook for this connection
"""
- connection = BaseHook.get_connection(self.conn_id)
- hook = connection.get_hook(hook_params=self.hook_params)
+ hook = await get_async_hook(self.conn_id, hook_params=self.hook_params)
if not isinstance(hook, DbApiHook):
- raise AirflowException(
- f"You are trying to use `common-sql` with
{hook.__class__.__name__},"
- " but its provider does not support it. Please upgrade the
provider to a version that"
- " supports `common-sql`. The hook class should be a subclass
of"
- f" `{hook.__class__.__module__}.{hook.__class__.__name__}`."
- f" Got {hook.__class__.__name__} hook with class hierarchy:
{hook.__class__.mro()}"
+ raise TypeError(
+ f"You are trying to use the SqlExecuteQueryOperator in
deferrable mode with {hook.__class__.__name__},"
+ " but its provider does not support this. Please set
deferrable=False"
+ f" Got {hook.__class__.__name__} with class hierarchy:
{hook.__class__.mro()}"
+ )
+ if self.read_only and not (hook.supports_async_execution() and
hook.supports_readonly_execution()):
+ raise NotImplementedError(
+ f"{hook.__class__.__name__} does not support read-only
execution, so it cannot run a"
+ " deferred query safely (a triggerer restart could re-run it).
Set"
+ " enforce_read_only=False to run without the read-only guard
if the query is"
+ " idempotent, or deferrable=False to run it on the worker."
)
return hook
- async def _get_records(self) -> Any:
- from asgiref.sync import sync_to_async
-
- hook = self.get_hook()
+ async def _run_query(self, hook: DbApiHook, fetch_results: bool):
+ """
+ Run the query against ``hook``, using its native async driver if
available.
+ Hooks without a real async driver
(:meth:`DbApiHook.supports_async_execution`) fall back to
+ running the synchronous :meth:`DbApiHook.run` in a worker thread,
matching the compatibility the
+ `GenericTransfer` operator has always relied on for arbitrary
DB-specific hooks.
+ """
+ if hook.supports_async_execution():
+ if fetch_results:
+ return await hook.arun(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=fetch_all_handler,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ read_only=self.read_only,
+ )
+ return await hook.arun(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=None,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ read_only=self.read_only,
+ )
if AIRFLOW_V_3_2_PLUS:
# This is only supported from Airflow 3.2 or higher due to added
async support in CommsDecoder
- return await sync_to_async(hook.get_records)(self.sql)
- return hook.get_records(self.sql)
+ # `sync_to_async` erases `run`'s overloads, so it is cast to a
plain callable first.
+ arun_in_thread = sync_to_async(cast("Callable[..., Any]",
hook.run))
+ if fetch_results:
+ return await arun_in_thread(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=fetch_all_handler,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ )
+ return await arun_in_thread(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=None,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ )
+ if fetch_results:
+ return hook.run(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=fetch_all_handler,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ )
+ return hook.run(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=None,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ )
async def run(self) -> AsyncIterator[TriggerEvent]:
try:
+ hook = await self.aget_hook()
+
self.log.info("Extracting data from %s", self.conn_id)
self.log.info("Executing: \n %s", self.sql)
- self.log.info("Reading records from %s", self.conn_id)
- results = await self._get_records()
+ 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 self._run_query(hook, fetch_results=True)
+
+ self.log.info("Executing query from %s done!", self.conn_id)
+ self.log.debug("results: %s", results)
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "results": results,
Review Comment:
`descriptions` get normalized via `_jsonsafe_descriptions`, but the
`results` rows are returned raw. A Postgres row holding platform-dependant
types (e.g. timestamptz) has to survive TriggerEvent serialization and the trip
back to the worker.
Could you add a test that selects numeric/timestamptz/bytea/uuid and asserts
the round-tripped values match `deferrable=False`? If serde doesn't cover a
type we'd silently corrupt or fail results in the exact path this PR adds. (I
honestly do not know what we did around here before this change.)
--
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]
