ferruzzi commented on code in PR #70714:
URL: https://github.com/apache/airflow/pull/70714#discussion_r3731818643
##########
airflow-core/src/airflow/models/deadline.py:
##########
@@ -320,34 +343,18 @@ def get_reference_class(cls, reference_name: str) ->
type[BaseDeadlineReference]
class BaseDeadlineReference(LoggingMixin, ABC):
"""Base class for all Deadline implementations."""
- # Set of required kwargs - subclasses should override this.
- required_kwargs: set[str] = set()
-
@classproperty
def reference_name(cls: Any) -> str:
return cls.__name__
def evaluate_with(self, *, session: Session, interval: timedelta,
**kwargs: Any) -> datetime | None:
- """Validate the provided kwargs and evaluate this deadline with
the given conditions."""
- filtered_kwargs = {k: v for k, v in kwargs.items() if k in
self.required_kwargs}
-
- if missing_kwargs := self.required_kwargs - filtered_kwargs.keys():
- raise ValueError(
- f"{self.__class__.__name__} is missing required
parameters: {', '.join(missing_kwargs)}"
- )
-
- if extra_kwargs := kwargs.keys() - filtered_kwargs.keys():
- self.log.debug(
- "%s ignoring unexpected parameters: %s",
- self.reference_name,
- ", ".join(extra_kwargs),
- )
-
- base_time = self._evaluate_with(session=session, **filtered_kwargs)
+ """Evaluate this deadline with the supplied context."""
+ evaluation_kwargs = _get_evaluation_kwargs(self,
self._evaluate_with, kwargs)
+ base_time = self._evaluate_with(session=session,
**evaluation_kwargs)
return base_time + interval if base_time is not None else None
@abstractmethod
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
+ def _evaluate_with(self, *, session: Session, dagrun: Any) -> datetime
| None:
Review Comment:
Rather than using Any, I think a cleaner solution might be to type `dagrun`
as a `DagRunProtocol` and add `queued_at` to that protocol.
##########
airflow-core/src/airflow/models/deadline.py:
##########
@@ -381,7 +388,7 @@ class FixedDatetimeDeadline(BaseDeadlineReference):
_datetime: datetime
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
+ def _evaluate_with(self, *, session: Session, dagrun: Any) -> datetime
| None:
Review Comment:
Rather than passing `dagrun` and not using it, I think we can widen the
contract here by replacing it with `**kwargs` in any case where dagrun isn't
actually used (this is the only one out of the built-in references):
```suggestion
def _evaluate_with(self, *, session: Session, **kwargs) -> datetime
| None:
```
##########
airflow-core/src/airflow/models/deadline.py:
##########
@@ -57,6 +59,27 @@
CALLBACK_METRICS_PREFIX = "deadline_alerts"
+def _get_evaluation_kwargs(reference: Any, evaluator: Any, kwargs: dict[str,
Any]) -> dict[str, Any]:
+ """Return the evaluation arguments accepted by a deadline reference."""
+ required_kwargs: set[str] | None = getattr(reference, "required_kwargs",
None)
+ if required_kwargs is not None:
+ warnings.warn(
+ "required_kwargs is deprecated. Declare the keyword-only
parameters your "
+ "_evaluate_with() implementation needs instead.",
+ RemovedInAirflow4Warning,
+ stacklevel=3,
+ )
+ kwargs = {key: value for key, value in kwargs.items() if key in
required_kwargs}
+ if missing_kwargs := required_kwargs - kwargs.keys():
+ raise ValueError(
+ f"{reference.__class__.__name__} is missing required
parameters: {', '.join(missing_kwargs)}"
+ )
+ return kwargs
+
+ parameters = signature(evaluator).parameters
+ return {key: value for key, value in kwargs.items() if key in parameters}
Review Comment:
With the way you wrote it, if the user's custom `_evaluate_with` accepts
`**kwargs`, they would all get dropped when they very likely intend kwargs to
mean "give me anything", which is a standard Python understanding. This way if
`_evaluate_with` claims to accept anything, then give it everything:
```suggestion
parameters = signature(evaluator).parameters
if any(param.kind is param.VAR_KEYWORD for param in parameters.values()):
return kwargs
return {key: value for key, value in kwargs.items() if key in parameters}
```
It might mean changing some of your tests as well;
`test_legacy_serialized_custom_reference_ignores_evaluation_context` for sure
will need an adjustment in the asserts.
##########
airflow-core/docs/howto/deadline-alerts.rst:
##########
@@ -441,24 +441,17 @@ choose a different time.
class MyCustomDecoratedReference(BaseDeadlineReference):
"""A custom reference evaluated when Dag runs are created."""
- def _evaluate_with(self, *, session: Session, **kwargs) -> datetime:
- # Add your business logic here
- return your_datetime
+ def _evaluate_with(self, *, session: Session, dagrun) -> datetime:
+ return dagrun.logical_date
Review Comment:
For these two examples and the two down in the `deadline_reference`
docstring, can we add back the `# Add your business logic here` comments and
not return the dagrun.logical_date directly? We should be showing that
dagrun.logical_date is available to use in their business logic, not that it is
or will be returned.
Maybe something like this gets the point across?
```suggestion
def _evaluate_with(self, *, session: Session, dagrun) -> datetime:
my_datetime = my_business_logic(dagrun.logical_date)
return my_datetime
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
