# 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 typing
from collections.abc import Collection
import attrs
from airflow.exceptions import AirflowTimetableInvalid
from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetAll, SerializedAssetBase
from airflow.timetables.base import Timetable
from airflow.timetables.simple import (
AssetTriggeredTimetable,
_coerce_asset_condition,
_get_default_batch_asset_events,
)
from airflow.utils.types import DagRunType
if typing.TYPE_CHECKING:
import pendulum
from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction
def _validate_asset_time_schedule(*, timetable: Timetable, asset_condition: SerializedAssetBase) -> None:
if timetable.asset_triggered or timetable.asset_gated:
raise AirflowTimetableInvalid("Cannot nest asset-aware timetables")
if not isinstance(asset_condition, SerializedAssetBase):
raise AirflowTimetableInvalid("All elements in 'assets' must be assets")
@attrs.define(eq=False, repr=False, kw_only=True, slots=False)
[docs]
class AssetOrTimeSchedule(AssetTriggeredTimetable):
"""Combine time-based scheduling with event-based scheduling."""
[docs]
asset_condition: SerializedAssetBase = attrs.field(alias="assets", converter=_coerce_asset_condition)
[docs]
batch_asset_events: bool = attrs.field(
default=attrs.Factory(_get_default_batch_asset_events, takes_self=True)
)
[docs]
def __attrs_post_init__(self) -> None:
self.description = f"Triggered by assets or {self.timetable.description}"
self.periodic = self.timetable.periodic
self.can_be_scheduled = self.timetable.can_be_scheduled
self.active_runs_limit = self.timetable.active_runs_limit
@classmethod
[docs]
def deserialize(cls, data: dict[str, typing.Any]) -> Timetable:
from airflow.serialization.decoders import decode_asset_like, decode_timetable
return cls(
assets=decode_asset_like(data["asset_condition"]),
timetable=decode_timetable(data["timetable"]),
batch_asset_events=data.get("batch_asset_events", True),
)
[docs]
def validate(self) -> None:
super().validate()
_validate_asset_time_schedule(
timetable=self.timetable,
asset_condition=self.asset_condition,
)
[docs]
def serialize(self) -> dict[str, typing.Any]:
from airflow.serialization.encoders import encode_asset_like, encode_timetable
return {
"asset_condition": encode_asset_like(self.asset_condition),
"timetable": encode_timetable(self.timetable),
"batch_asset_events": self.batch_asset_events,
}
@property
[docs]
def summary(self) -> str:
return f"Asset or {self.timetable.summary}"
[docs]
def infer_manual_data_interval(self, *, run_after: pendulum.DateTime) -> DataInterval:
return self.timetable.infer_manual_data_interval(run_after=run_after)
[docs]
def next_dagrun_info(
self, *, last_automated_data_interval: DataInterval | None, restriction: TimeRestriction
) -> DagRunInfo | None:
return self.timetable.next_dagrun_info(
last_automated_data_interval=last_automated_data_interval,
restriction=restriction,
)
[docs]
def generate_run_id(self, *, run_type: DagRunType, **kwargs: typing.Any) -> str:
if run_type != DagRunType.ASSET_TRIGGERED:
return self.timetable.generate_run_id(run_type=run_type, **kwargs)
return super().generate_run_id(run_type=run_type, **kwargs)
[docs]
class AssetAndTimeSchedule(Timetable):
"""
Time-based schedule that waits for required assets before creating a run.
This timetable composes a time-based timetable with an asset condition. It
schedules runs according to the provided ``timetable`` (e.g. cron), but a run
is only created when all required assets are present. Unlike
:class:`AssetOrTimeSchedule`, this does not create asset-triggered runs.
"""
def __init__(
self,
*,
timetable: Timetable,
assets: Collection[SerializedAsset] | SerializedAssetBase,
) -> None:
from airflow.serialization.encoders import ensure_serialized_asset
if isinstance(assets, SerializedAssetBase):
self.asset_condition = assets
elif isinstance(assets, Collection):
self.asset_condition = SerializedAssetAll([ensure_serialized_asset(a) for a in assets])
else:
self.asset_condition = ensure_serialized_asset(assets)
@classmethod
[docs]
def deserialize(cls, data: dict[str, typing.Any]) -> Timetable:
from airflow.serialization.decoders import decode_asset_like, decode_timetable
return cls(
assets=decode_asset_like(data["asset_condition"]),
timetable=decode_timetable(data["timetable"]),
)
[docs]
def serialize(self) -> dict[str, typing.Any]:
from airflow.serialization.encoders import encode_asset_like, encode_timetable
return {
"asset_condition": encode_asset_like(self.asset_condition),
"timetable": encode_timetable(self.timetable),
}
[docs]
def validate(self) -> None:
_validate_asset_time_schedule(
timetable=self.timetable,
asset_condition=self.asset_condition,
)
@property
[docs]
def description(self) -> str: # type: ignore[override]
return f"Triggered by assets and {self.timetable.description}"
@property
[docs]
def summary(self) -> str:
return f"Asset and {self.timetable.summary}"
@property
[docs]
def periodic(self) -> bool: # type: ignore[override]
return self.timetable.periodic
@property
[docs]
def can_be_scheduled(self) -> bool: # type: ignore[override]
return self.timetable.can_be_scheduled
@property
[docs]
def active_runs_limit(self) -> int | None: # type: ignore[override]
return self.timetable.active_runs_limit
[docs]
def infer_manual_data_interval(self, *, run_after: pendulum.DateTime) -> DataInterval:
return self.timetable.infer_manual_data_interval(run_after=run_after)
[docs]
def next_dagrun_info(
self, *, last_automated_data_interval: DataInterval | None, restriction: TimeRestriction
) -> DagRunInfo | None:
return self.timetable.next_dagrun_info(
last_automated_data_interval=last_automated_data_interval,
restriction=restriction,
)
[docs]
def generate_run_id(self, *, run_type: DagRunType, **kwargs: typing.Any) -> str:
# All run IDs are delegated to the wrapped timetable; this class
# intentionally does not create ASSET_TRIGGERED runs.
return self.timetable.generate_run_id(run_type=run_type, **kwargs)