Add setters to MappedOperator on_*_callbacks (apache#28313)

sekikn · Dec 13, 2022 · 105dbd2 · 105dbd2
1 parent 02c2b2c
commit 105dbd2
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
@@ -450,18 +450,34 @@ def resources(self) -> Resources | None:
     def on_execute_callback(self) -> TaskStateChangeCallback | None:
         return self.partial_kwargs.get("on_execute_callback")
 
+    @on_execute_callback.setter
+    def on_execute_callback(self, value: TaskStateChangeCallback | None) -> None:
+        self.partial_kwargs["on_execute_callback"] = value
+
     @property
     def on_failure_callback(self) -> TaskStateChangeCallback | None:
         return self.partial_kwargs.get("on_failure_callback")
 
+    @on_failure_callback.setter
+    def on_failure_callback(self, value: TaskStateChangeCallback | None) -> None:
+        self.partial_kwargs["on_failure_callback"] = value
+
     @property
     def on_retry_callback(self) -> TaskStateChangeCallback | None:
         return self.partial_kwargs.get("on_retry_callback")
 
+    @on_retry_callback.setter
+    def on_retry_callback(self, value: TaskStateChangeCallback | None) -> None:
+        self.partial_kwargs["on_retry_callback"] = value
+
     @property
     def on_success_callback(self) -> TaskStateChangeCallback | None:
         return self.partial_kwargs.get("on_success_callback")
 
+    @on_success_callback.setter
+    def on_success_callback(self, value: TaskStateChangeCallback | None) -> None:
+        self.partial_kwargs["on_success_callback"] = value
+
     @property
     def run_as_user(self) -> str | None:
         return self.partial_kwargs.get("run_as_user")