Source code for airflow_pydantic.sensors.timedelta

from datetime import timedelta
from logging import getLogger

from pydantic import Field, field_validator

from ..core import Task
from ..utils import ImportPath
from .base import BaseSensorArgs

__all__ = (
    "TimeDeltaSensor",
    "TimeDeltaSensorArgs",
    "WaitSensor",
    "WaitSensorArgs",
)

_log = getLogger(__name__)




class TimeDeltaSensorArgs(BaseSensorArgs):
    # timedelta sensor args
    # https://airflow.apache.org/docs/apache-airflow-providers-standard/stable/_api/airflow/providers/standard/sensors/time_delta/index.html#airflow.providers.standard.sensors.time_delta.TimeDeltaSensor
    delta: timedelta = Field(description="Time to wait before succeeding.")
    deferrable: bool | None = Field(default=None, description="If True, the sensor will operate in deferrable mode")





class TimeDeltaSensor(Task, TimeDeltaSensorArgs):
    operator: ImportPath = Field(default="airflow_pydantic.airflow.TimeDeltaSensor", description="airflow sensor path", validate_default=True)

    @field_validator("operator")
    @classmethod
    def validate_operator(cls, v: type) -> type:
        from airflow_pydantic.airflow import TimeDeltaSensor, _AirflowPydanticMarker

        if not isinstance(v, type):
            raise TypeError(f"operator must be 'airflow.providers.standard.sensors.time.TimeDeltaSensor', got: {v}")
        if issubclass(v, _AirflowPydanticMarker):
            _log.info("TimeSensor is a marker class, returning as is")
            return v
        if not issubclass(v, TimeDeltaSensor):
            raise TypeError(f"operator must be 'airflow.providers.standard.sensors.time.TimeDeltaSensor', got: {v}")
        return v





class WaitSensorArgs(BaseSensorArgs):
    # wait sensor args
    # https://airflow.apache.org/docs/apache-airflow-providers-standard/stable/_api/airflow/providers/standard/sensors/time_delta/index.html#airflow.providers.standard.sensors.time_delta.WaitSensor
    time_to_wait: timedelta | int = Field(description="Time length to wait after the task starts before succeeding.")
    deferrable: bool | None = Field(default=None, description="If True, the sensor will operate in deferrable mode")





class WaitSensor(Task, WaitSensorArgs):
    operator: ImportPath = Field(default="airflow_pydantic.airflow.WaitSensor", description="airflow sensor path", validate_default=True)

    @field_validator("operator")
    @classmethod
    def validate_operator(cls, v: type) -> type:
        from airflow_pydantic.airflow import WaitSensor, _AirflowPydanticMarker

        if not isinstance(v, type):
            raise TypeError(f"operator must be 'airflow.providers.standard.sensors.time.WaitSensor', got: {v}")
        if issubclass(v, _AirflowPydanticMarker):
            _log.info("WaitSensor is a marker class, returning as is")
            return v
        if not issubclass(v, WaitSensor):
            raise TypeError(f"operator must be 'airflow.providers.standard.sensors.time.WaitSensor', got: {v}")
        return v