feat: support include_deferred in pool sync for airflow 2.7.0 (#775)

* fix: add 'include_deferred' column to sync_pools script

Signed-off-by: Einav Daniel <[email protected]>

* fix: include_deferred implementation

Signed-off-by: Mathew Wicks <[email protected]>

---------

Signed-off-by: Einav Daniel <[email protected]>
Signed-off-by: Mathew Wicks <[email protected]>
Co-authored-by: Mathew Wicks <[email protected]>

Author: EinavDanielDX

charts/airflow/docs/faq/dags/airflow-pools.md

@@ -14,9 +14,12 @@ airflow:
     - name: "pool_1"
       description: "example pool with 5 slots"
       slots: 5
+      
     - name: "pool_2"
       description: "example pool with 10 slots"
       slots: 10
+      ## if deferred tasks count towards the slot limit, requires airflow 2.7.0+ (default: false)
+      #include_deferred: false
 
   ## if we create a Deployment to perpetually sync `airflow.pools`
   poolsUpdate: true

charts/airflow/templates/sync/_helpers/sync_pools.tpl

@@ -47,18 +47,23 @@ class PoolWrapper(object):
             name: str,
             description: str,
             slots: int,
+            include_deferred: bool,
             policies: List[ScheduledPolicy],
             enable_policies: bool,
     ):
         self.name = name
         self.description = description
         self.slots = slots
+        self.include_deferred = include_deferred
         self.policies = policies
         self.enable_policies = enable_policies
 
     def as_pool(self) -> Pool:
         pool = Pool()
         pool.pool = self.name
+        # NOTE: include_deferred is only available in Airflow 2.7.0+
+        if hasattr(Pool, "include_deferred"):
+            pool.include_deferred = self.include_deferred
         if self._has_policies():
             most_recent_policy = self._most_recent_policy()
             pool.slots = most_recent_policy.slots
@@ -92,6 +97,15 @@ VAR__POOL_WRAPPERS = {
     {{ required "the `slots` in each `airflow.pools[]` must be int-type!" nil }}
     {{- end }}
     slots={{ (required "the `slots` in each `airflow.pools[]` must be non-empty!" .slots) }},
+    {{- $include_deferred := dig "include_deferred" nil . }}
+    {{- if not (or (typeIs "bool" $include_deferred) (eq $include_deferred nil)) }}
+    {{ required "if specified, the `include_deferred` in each `airflow.pools[]` must be bool-type!" nil }}
+    {{- end }}
+    {{- if $include_deferred }}
+    include_deferred=True,
+    {{- else }}
+    include_deferred=False,
+    {{- end }}
     policies=[
         {{- range .policies }}
             ScheduledPolicy(
@@ -126,6 +140,7 @@ def compare_pools(p1: Pool, p2: Pool) -> bool:
             p1.pool == p1.pool
             and p1.description == p2.description
             and p1.slots == p2.slots
+            and getattr(p1, "include_deferred", False) == getattr(p2, "include_deferred", False)
     )
 
 
@@ -152,6 +167,8 @@ def sync_pool(pool_wrapper: PoolWrapper) -> None:
                 logging.info(f"Pool=`{p_name}` exists but has changed, updating...")
                 p_old.description = p_new.description
                 p_old.slots = p_new.slots
+                if hasattr(Pool, "include_deferred"):
+                    p_old.include_deferred = p_new.include_deferred
                 pool_updated = True
 
     if pool_added:

charts/airflow/values.yaml

@@ -162,6 +162,8 @@ airflow:
   ##     - name: "pool_2"
   ##       description: "example pool with 2 cron policies"
   ##       slots: 0
+  ##       ## if deferred tasks count towards the slot limit, requires airflow 2.7.0+ (default: false)
+  ##       include_deferred: false
   ##       ## at each sync interval, the policy with the most recently past `recurrence` is applied
   ##       policies:
   ##         - name: "scale up at 7pm UTC"