[WIP] Add a public interface for custom weight_rule implementation

airflow/example_dags/example_priority_weight_strategy.py

@@ -25,7 +25,6 @@
 
 from airflow.models.dag import DAG
 from airflow.operators.python import PythonOperator
-from airflow.task.priority_strategy import PriorityWeightStrategy
 
 if TYPE_CHECKING:
     from airflow.models import TaskInstance
@@ -36,13 +35,6 @@ def success_on_third_attempt(ti: TaskInstance, **context):
         raise Exception("Not yet")
 
 
-class DecreasingPriorityStrategy(PriorityWeightStrategy):
-    """A priority weight strategy that decreases the priority weight with each attempt."""
-
-    def get_weight(self, ti: TaskInstance):
-        return max(3 - ti._try_number + 1, 1)
-
-
 with DAG(
     dag_id="example_priority_weight_strategy",
     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
@@ -63,7 +55,5 @@ def get_weight(self, ti: TaskInstance):
     decreasing_weight_task = PythonOperator(
         task_id="decreasing_weight_task",
         python_callable=success_on_third_attempt,
-        priority_weight_strategy=(
-            "airflow.example_dags.example_priority_weight_strategy.DecreasingPriorityStrategy"
-        ),
+        priority_weight_strategy=("decreasing_priority_weight_strategy.DecreasingPriorityStrategy"),
     )

airflow/example_dags/plugins/decreasing_priority_weight_strategy.py

@@ -0,0 +1,37 @@
+# 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
+
+from typing import TYPE_CHECKING
+
+from airflow.plugins_manager import AirflowPlugin
+from airflow.task.priority_strategy import PriorityWeightStrategy
+
+if TYPE_CHECKING:
+    from airflow.models import TaskInstance
+
+
+class DecreasingPriorityStrategy(PriorityWeightStrategy):
+    """A priority weight strategy that decreases the priority weight with each attempt."""
+
+    def get_weight(self, ti: TaskInstance):
+        return max(3 - ti._try_number + 1, 1)
+
+
+class DecreasingPriorityWeightStrategyPlugin(AirflowPlugin):
+    name = "decreasing_priority_weight_strategy_plugin"
+    priority_weight_strategies = [DecreasingPriorityStrategy]

airflow/migrations/versions/0132_2_8_0_add_priority_weight_strategy_to_task_.py

@@ -26,7 +26,7 @@
 
 import sqlalchemy as sa
 from alembic import op
-
+from sqlalchemy import text
 
 # revision identifiers, used by Alembic.
 revision = "624ecf3b6a5e"
@@ -37,9 +37,18 @@
 
 
 def upgrade():
+    json_type = sa.JSON
+    conn = op.get_bind()
+    if conn.dialect.name != "postgresql":
+        # Mysql 5.7+/MariaDB 10.2.3 has JSON support. Rather than checking for
+        # versions, check for the function existing.
+        try:
+            conn.execute(text("SELECT JSON_VALID(1)")).fetchone()
+        except (sa.exc.OperationalError, sa.exc.ProgrammingError):
+            json_type = sa.Text
     """Apply add priority_weight_strategy to task_instance"""
     with op.batch_alter_table("task_instance") as batch_op:
-        batch_op.add_column(sa.Column("priority_weight_strategy", sa.String(length=1000)))
+        batch_op.add_column(sa.Column("priority_weight_strategy", json_type()))
 
 
 def downgrade():

airflow/models/abstractoperator.py

@@ -21,7 +21,7 @@
 import inspect
 import warnings
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Collection, Iterable, Iterator, Sequence
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Collection, Iterable, Iterator, Literal, Sequence
 
 from sqlalchemy import select
 
@@ -55,6 +55,7 @@
     from airflow.models.mappedoperator import MappedOperator
     from airflow.models.operator import Operator
     from airflow.models.taskinstance import TaskInstance
+    from airflow.task.priority_strategy import PriorityWeightStrategy
     from airflow.utils.task_group import TaskGroup
 
 DEFAULT_OWNER: str = conf.get_mandatory_value("operators", "default_owner")
@@ -106,7 +107,7 @@ class AbstractOperator(Templater, DAGNode):
     operator_class: type[BaseOperator] | dict[str, Any]
 
     weight_rule: str | None
-    priority_weight_strategy: str
+    priority_weight_strategy: Literal["absolute", "downstream", "upstream"] | PriorityWeightStrategy
     priority_weight: int
 
     # Defines the operator level extra links.
@@ -206,6 +207,23 @@ def on_failure_fail_dagrun(self, value):
             )
         self._on_failure_fail_dagrun = value
 
+    @property
+    def parsed_priority_weight_strategy(self) -> PriorityWeightStrategy:
+        from airflow.serialization.serialized_objects import _get_registered_priority_weight_strategy
+        from airflow.utils.module_loading import qualname
+
+        if isinstance(self.priority_weight_strategy, str):
+            priority_weight_strategy_cls = _get_registered_priority_weight_strategy(
+                self.priority_weight_strategy
+            )
+            if priority_weight_strategy_cls is None:
+                raise AirflowException(f"Unknown priority strategy {priority_weight_strategy_cls}")
+            return priority_weight_strategy_cls()
+        priority_weight_strategy_str = qualname(self.priority_weight_strategy)
+        if _get_registered_priority_weight_strategy(priority_weight_strategy_str) is None:
+            raise AirflowException(f"Unknown priority strategy {priority_weight_strategy_str}")
+        return self.priority_weight_strategy
+
     def as_setup(self):
         self.is_setup = True
         return self

airflow/models/baseoperator.py

@@ -77,7 +77,6 @@
 from airflow.models.taskinstance import TaskInstance, clear_task_instances
 from airflow.models.taskmixin import DependencyMixin
 from airflow.serialization.enums import DagAttributeTypes
-from airflow.task.priority_strategy import get_priority_weight_strategy
 from airflow.ti_deps.deps.not_in_retry_period_dep import NotInRetryPeriodDep
 from airflow.ti_deps.deps.not_previously_skipped_dep import NotPreviouslySkippedDep
 from airflow.ti_deps.deps.prev_dagrun_dep import PrevDagrunDep
@@ -796,6 +795,8 @@ def __init__(
         **kwargs,
     ):
         from airflow.models.dag import DagContext
+        from airflow.serialization.serialized_objects import _get_registered_priority_weight_strategy
+        from airflow.utils.module_loading import qualname
         from airflow.utils.task_group import TaskGroupContext
 
         self.__init_kwargs = {}
@@ -921,7 +922,15 @@ def __init__(
             )
             self.priority_weight_strategy = weight_rule
         # validate the priority weight strategy
-        get_priority_weight_strategy(self.priority_weight_strategy)
+        # validate the priority weight strategy
+        priority_weight_strategy_cls = (
+            self.priority_weight_strategy
+            if isinstance(self.priority_weight_strategy, str)
+            else qualname(self.priority_weight_strategy)
+        )
+        if _get_registered_priority_weight_strategy(priority_weight_strategy_cls) is None:
+            raise AirflowException(f"Unknown priority strategy {priority_weight_strategy_cls}")
+
         self.resources = coerce_resources(resources)
         if task_concurrency and not max_active_tis_per_dag:
             # TODO: Remove in Airflow 3.0

airflow/models/mappedoperator.py

@@ -49,7 +49,6 @@
 )
 from airflow.models.pool import Pool
 from airflow.serialization.enums import DagAttributeTypes
-from airflow.task.priority_strategy import get_priority_weight_strategy
 from airflow.ti_deps.deps.mapped_task_expanded import MappedTaskIsExpanded
 from airflow.typing_compat import Literal
 from airflow.utils.context import context_update_for_unmapped
@@ -79,6 +78,7 @@
     from airflow.models.operator import Operator
     from airflow.models.param import ParamsDict
     from airflow.models.xcom_arg import XComArg
+    from airflow.task.priority_strategy import PriorityWeightStrategy
     from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
     from airflow.utils.context import Context
     from airflow.utils.operator_resources import Resources
@@ -314,6 +314,8 @@ def __repr__(self):
 
     def __attrs_post_init__(self):
         from airflow.models.xcom_arg import XComArg
+        from airflow.serialization.serialized_objects import _get_registered_priority_weight_strategy
+        from airflow.utils.module_loading import qualname
 
         if self.get_closest_mapped_task_group() is not None:
             raise NotImplementedError("operator expansion in an expanded task group is not yet supported")
@@ -332,7 +334,13 @@ def __attrs_post_init__(self):
                 f"{self.task_id!r}."
             )
         # validate the priority weight strategy
-        get_priority_weight_strategy(self.priority_weight_strategy)
+        priority_weight_strategy_cls = (
+            self.priority_weight_strategy
+            if isinstance(self.priority_weight_strategy, str)
+            else qualname(self.priority_weight_strategy)
+        )
+        if _get_registered_priority_weight_strategy(priority_weight_strategy_cls) is None:
+            raise AirflowException(f"Unknown priority strategy {priority_weight_strategy_cls}")
 
     @classmethod
     @cache
@@ -479,7 +487,7 @@ def weight_rule(self) -> str | None:  # type: ignore[override]
         return self.partial_kwargs.get("weight_rule") or DEFAULT_WEIGHT_RULE
 
     @property
-    def priority_weight_strategy(self) -> str:  # type: ignore[override]
+    def priority_weight_strategy(self) -> str | PriorityWeightStrategy:  # type: ignore[override]
         return (
             self.weight_rule  # for backward compatibility
             or self.partial_kwargs.get("priority_weight_strategy")

airflow/models/taskinstance.py

@@ -21,6 +21,7 @@
 import contextlib
 import hashlib
 import itertools
+import json
 import logging
 import math
 import operator
@@ -37,6 +38,7 @@
 import jinja2
 import lazy_object_proxy
 import pendulum
+import sqlalchemy_jsonfield
 from jinja2 import TemplateAssertionError, UndefinedError
 from sqlalchemy import (
     Column,
@@ -98,7 +100,6 @@
 from airflow.plugins_manager import integrate_macros_plugins
 from airflow.sentry import Sentry
 from airflow.stats import Stats
-from airflow.task.priority_strategy import get_priority_weight_strategy
 from airflow.templates import SandboxedEnvironment
 from airflow.ti_deps.dep_context import DepContext
 from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS
@@ -148,6 +149,7 @@
     from airflow.models.operator import Operator
     from airflow.serialization.pydantic.dag import DagModelPydantic
     from airflow.serialization.pydantic.taskinstance import TaskInstancePydantic
+    from airflow.task.priority_strategy import PriorityWeightStrategy
     from airflow.timetables.base import DataInterval
     from airflow.typing_compat import Literal, TypeGuard
     from airflow.utils.task_group import TaskGroup
@@ -883,9 +885,7 @@ def _refresh_from_task(
     task_instance.pool_slots = task.pool_slots
     with contextlib.suppress(Exception):
         # This method is called from the different places, and sometimes the TI is not fully initialized
-        task_instance.priority_weight = get_priority_weight_strategy(
-            task.priority_weight_strategy
-        ).get_weight(
+        task_instance.priority_weight = task.priority_weight_strategy.get_weight(
             task_instance  # type: ignore
         )
     task_instance.run_as_user = task.run_as_user
@@ -1222,7 +1222,7 @@ class TaskInstance(Base, LoggingMixin):
     pool_slots = Column(Integer, default=1, nullable=False)
     queue = Column(String(256))
     priority_weight = Column(Integer)
-    priority_weight_strategy = Column(String(1000))
+    priority_weight_strategy = Column(sqlalchemy_jsonfield.JSONField(json=json))
     operator = Column(String(1000))
     custom_operator_name = Column(String(1000))
     queued_dttm = Column(UtcDateTime)
@@ -1391,7 +1391,9 @@ def insert_mapping(run_id: str, task: Operator, map_index: int) -> dict[str, Any
 
         :meta private:
         """
-        priority_weight = get_priority_weight_strategy(task.priority_weight_strategy).get_weight(
+        from airflow.serialization.serialized_objects import _encode_priority_weight_strategy
+
+        priority_weight = task.parsed_priority_weight_strategy.get_weight(
             TaskInstance(task=task, run_id=run_id, map_index=map_index)
         )
         return {
@@ -1405,7 +1407,9 @@ def insert_mapping(run_id: str, task: Operator, map_index: int) -> dict[str, Any
             "pool": task.pool,
             "pool_slots": task.pool_slots,
             "priority_weight": priority_weight,
-            "priority_weight_strategy": task.priority_weight_strategy,
+            "priority_weight_strategy": _encode_priority_weight_strategy(
+                task.parsed_priority_weight_strategy
+            ),
             "run_as_user": task.run_as_user,
             "max_tries": task.retries,
             "executor_config": task.executor_config,
@@ -3462,7 +3466,7 @@ def __init__(
         key: TaskInstanceKey,
         run_as_user: str | None = None,
         priority_weight: int | None = None,
-        priority_weight_strategy: str | None = None,
+        priority_weight_strategy: PriorityWeightStrategy | None = None,
     ):
         self.dag_id = dag_id
         self.task_id = task_id
@@ -3502,6 +3506,8 @@ def as_dict(self):
 
     @classmethod
     def from_ti(cls, ti: TaskInstance) -> SimpleTaskInstance:
+        from airflow.serialization.serialized_objects import _decode_priority_weight_strategy
+
         return cls(
             dag_id=ti.dag_id,
             task_id=ti.task_id,
@@ -3517,7 +3523,7 @@ def from_ti(cls, ti: TaskInstance) -> SimpleTaskInstance:
             key=ti.key,
             run_as_user=ti.run_as_user if hasattr(ti, "run_as_user") else None,
             priority_weight=ti.priority_weight if hasattr(ti, "priority_weight") else None,
-            priority_weight_strategy=ti.priority_weight_strategy,
+            priority_weight_strategy=_decode_priority_weight_strategy(ti.priority_weight_strategy),
         )
 
     @classmethod