karenbraganz commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r4354014405
##########
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:
I ran some tests and found that bytes serialization fails because serde
doesn't cover it. I will raise an error if this data type is returned.
--
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]
