Source code for airflow.providers.common.sql.triggers.sql
#
# 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
from typing import TYPE_CHECKING
from airflow.exceptions import AirflowException
from airflow.hooks.base import BaseHook
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
[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
"""
def __init__(
self,
sql: str | list[str],
conn_id: str,
hook_params: dict | None = None,
**kwargs,
):
super().__init__(**kwargs)
[docs]
self.hook_params = hook_params
[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,
},
)
[docs]
async def run(self) -> AsyncIterator[TriggerEvent]:
try:
hook = BaseHook.get_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()}"
)
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 = hook.get_records(self.sql)
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)})