#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import asyncio
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.sql import DbApiHook
from airflow.triggers.base import BaseTrigger, TriggerEvent
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterable, Mapping
from typing import Any
[docs]
NO_ROWS_MESSAGE = "No rows returned, raising as per fail_on_empty flag"
[docs]
class SQLExecuteQueryTrigger(BaseTrigger):
"""
A trigger that executes SQL code in async mode.
: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 hook_params: hook parameters
:param parameters: The parameters to render the SQL query with (optional).
"""
def __init__(
self,
sql: str | list[str],
conn_id: str,
hook_params: dict | None = None,
parameters: Iterable | Mapping[str, Any] | None = None,
**kwargs,
):
super().__init__(**kwargs)
[docs]
self.hook_params = hook_params
[docs]
self.parameters = parameters
[docs]
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,
"hook_params": self.hook_params,
"parameters": self.parameters,
},
)
[docs]
def get_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)
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()}"
)
return hook
def _fetch_records(self, hook: DbApiHook) -> Any:
if self.parameters is None:
return hook.get_records(self.sql)
return hook.get_records(self.sql, self.parameters)
async def _get_records(self) -> Any:
hook = self.get_hook()
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(self._fetch_records)(hook)
return self._fetch_records(hook)
[docs]
async def run(self) -> AsyncIterator[TriggerEvent]:
try:
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()
self.log.info("Reading records from %s done!", self.conn_id)
self.log.debug("results: %s", results)
yield TriggerEvent({"status": "success", "results": results})
except Exception as e:
self.log.exception("An error occurred: %s", e)
yield TriggerEvent({"status": "failure", "message": str(e)})
[docs]
class SqlSensorTrigger(SQLExecuteQueryTrigger):
"""
Poll a SQL query until the first cell of the first row is truthy.
This is the default ``SqlSensor`` condition. The trigger waits ``poke_interval``
seconds before each query, because the sensor checks once on the worker before
deferring. Each poll, connection lookup included, runs in a thread so a slow query
does not block other triggers. The triggerer only serves connection lookups from
threads safely on Airflow 2 and on 3.2.2 or later, so ``SqlSensor`` only defers there.
An error during a poll ends the wait with a failure event, so the task fails and
retries as it would if the same error happened in poke mode.
:param sql: The SQL query to poll.
:param conn_id: The connection ID used to connect to the database.
:param poke_interval: Seconds to wait before each query.
:param parameters: The parameters to render the SQL query with (optional).
:param hook_params: Extra config params to be passed to the underlying hook.
:param fail_on_empty: Yield a failure event when the query returns no rows.
"""
def __init__(
self,
*,
sql: str,
conn_id: str,
poke_interval: float,
parameters: Iterable | Mapping[str, Any] | None = None,
hook_params: dict | None = None,
fail_on_empty: bool = False,
**kwargs,
):
super().__init__(sql=sql, conn_id=conn_id, hook_params=hook_params, parameters=parameters, **kwargs)
[docs]
self.poke_interval = poke_interval
[docs]
self.fail_on_empty = fail_on_empty
[docs]
def serialize(self) -> tuple[str, dict[str, Any]]:
"""Serialize the SqlSensorTrigger arguments and classpath."""
classpath, kwargs = super().serialize()
return classpath, {**kwargs, "poke_interval": self.poke_interval, "fail_on_empty": self.fail_on_empty}
def _poll(self) -> Any:
return self._fetch_records(self.get_hook())
[docs]
async def run(self) -> AsyncIterator[TriggerEvent]:
try:
while True:
await asyncio.sleep(self.poke_interval)
self.log.info("Poking: %s (with parameters %s)", self.sql, self.parameters)
records = await sync_to_async(self._poll, thread_sensitive=False)()
if not records:
if self.fail_on_empty:
yield TriggerEvent({"status": "failure", "message": NO_ROWS_MESSAGE})
return
elif records[0][0]:
yield TriggerEvent({"status": "success"})
return
self.log.info("Criteria not met; checking again in %s seconds.", self.poke_interval)
except Exception as e:
self.log.exception("An error occurred: %s", e)
yield TriggerEvent({"status": "failure", "message": str(e)})