diff --git a/astronomer/providers/snowflake/example_dags/example_snowflake_sensor.py b/astronomer/providers/snowflake/example_dags/example_snowflake_sensor.py index 6a9717a52..53ba7fc2c 100644 --- a/astronomer/providers/snowflake/example_dags/example_snowflake_sensor.py +++ b/astronomer/providers/snowflake/example_dags/example_snowflake_sensor.py @@ -67,7 +67,7 @@ sql=SNOWFLAKE_SLACK_SQL, poke_interval=POKE_INTERVAL, timeout=TASK_TIMEOUT * 60, - success=lambda result: True, # type: ignore[arg-type] + success=lambda result: True, ) snowflake_op_sql_str >> snowflake_op_with_params >> snowflake_op_sql_sensor >> snowflake_with_callable diff --git a/astronomer/providers/snowflake/sensors/snowflake.py b/astronomer/providers/snowflake/sensors/snowflake.py index b40d0a743..f5f5106b9 100644 --- a/astronomer/providers/snowflake/sensors/snowflake.py +++ b/astronomer/providers/snowflake/sensors/snowflake.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import timedelta -from typing import Any, Sequence +from typing import Any, Callable, Sequence from airflow.exceptions import AirflowException from airflow.providers.common.sql.sensors.sql import SqlSensor @@ -48,8 +48,8 @@ def __init__( snowflake_conn_id: str, sql: str, parameters: dict[str, Any] | None = None, - success: str | None = None, - failure: str | None = None, + success: Callable[[Any], bool] | None = None, + failure: Callable[[Any], bool] | None = None, fail_on_empty: bool = False, hook_params: dict[str, Any] | None = None, **kwargs: Any,