feat: add global task framework (#36368)

tests/integration_tests/tasks/test_throttling.py

@@ -0,0 +1,172 @@
+# 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.
+"""Integration tests for TaskContext update_task throttling.
+
+Tests verify:
+1. Final state is persisted correctly via cleanup flush
+2. Throttled updates are deferred, timer writes latest pending update
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+
+from superset_core.api.tasks import TaskScope, TaskStatus
+
+from superset.daos.tasks import TaskDAO
+from superset.extensions import db
+from superset.models.tasks import Task
+from superset.tasks.ambient_context import get_context
+from superset.tasks.registry import TaskRegistry
+from superset.tasks.scheduler import execute_task
+from tests.integration_tests.base_tests import SupersetTestCase
+from tests.integration_tests.constants import ADMIN_USERNAME
+
+
+def task_with_throttled_updates() -> None:
+    """Task with rapid progress and payload updates (exercises throttling)."""
+    ctx = get_context()
+
+    # Rapid-fire updates within throttle window
+    for i in range(10):
+        ctx.update_task(progress=(i + 1, 10), payload={"step": i + 1})
+
+
+def _register_test_tasks() -> None:
+    """Register test task functions if not already registered.
+
+    Called in setUp() to ensure tasks are registered regardless of
+    whether other tests have cleared the registry.
+    """
+    if not TaskRegistry.is_registered("test_throttle_combined"):
+        TaskRegistry.register("test_throttle_combined", task_with_throttled_updates)
+
+
+class TestUpdateTaskThrottling(SupersetTestCase):
+    """Integration test for update_task() throttling behavior."""
+
+    def setUp(self) -> None:
+        super().setUp()
+        self.login(ADMIN_USERNAME)
+        _register_test_tasks()
+
+    def test_throttled_updates_persisted_on_cleanup(self) -> None:
+        """Final state should be persisted regardless of throttling.
+
+        Verifies the core invariant: cleanup flush ensures final state is persisted.
+        """
+        task_obj = TaskDAO.create_task(
+            task_type="test_throttle_combined",
+            task_key=f"test_key_{uuid.uuid4().hex[:8]}",
+            task_name="Test Throttled Updates",
+            scope=TaskScope.SYSTEM,
+        )
+
+        # Use str(uuid) since Celery serializes args as JSON strings
+        result = execute_task.apply(
+            args=[str(task_obj.uuid), "test_throttle_combined", (), {}]
+        )
+
+        assert result.successful()
+        assert result.result["status"] == TaskStatus.SUCCESS.value
+
+        # Verify final state is persisted
+        db.session.expire_all()
+        task_obj = db.session.query(Task).filter_by(uuid=task_obj.uuid).first()
+
+        # Progress: 10/10 = 100%
+        props = task_obj.properties_dict
+        assert props.get("progress_current") == 10
+        assert props.get("progress_total") == 10
+        assert props.get("progress_percent") == 1.0
+
+        # Payload: final step
+        payload = task_obj.payload_dict
+        assert payload.get("step") == 10
+
+    def test_throttle_behavior(self) -> None:
+        """Test complete throttle behavior: immediate write, deferral, and timer.
+
+        Verifies:
+        1. First update writes immediately
+        2. Second and third updates within throttle window are deferred
+        3. Deferred timer fires and writes the LATEST pending update (third)
+        """
+        from flask import current_app
+
+        from superset.commands.tasks.submit import SubmitTaskCommand
+        from superset.tasks.context import TaskContext
+
+        # Get throttle interval from config (default: 2 seconds)
+        throttle_interval = current_app.config["TASK_PROGRESS_UPDATE_THROTTLE_INTERVAL"]
+
+        # Create task
+        task_obj = SubmitTaskCommand(
+            data={
+                "task_type": "test_throttle_behavior",
+                "task_key": f"test_key_{uuid.uuid4().hex[:8]}",
+                "task_name": "Test Throttle Behavior",
+                "scope": TaskScope.SYSTEM,
+            }
+        ).run()
+        task_uuid = task_obj.uuid
+
+        # Get fresh task for context
+        fresh_task = TaskDAO.find_one_or_none(uuid=task_uuid, skip_base_filter=True)
+        assert fresh_task is not None
+        ctx = TaskContext(fresh_task)
+
+        try:
+            # === Step 1: First update - writes immediately ===
+            ctx.update_task(progress=0.1, payload={"step": 1})
+
+            db.session.expire_all()
+            task_step1 = TaskDAO.find_one_or_none(uuid=task_uuid, skip_base_filter=True)
+            assert task_step1 is not None
+            assert task_step1.properties_dict.get("progress_percent") == 0.1
+            assert task_step1.payload_dict.get("step") == 1
+
+            # === Step 2: Second update - deferred (within throttle window) ===
+            ctx.update_task(progress=0.5, payload={"step": 2})
+
+            # === Step 3: Third update - also deferred, overwrites second in cache ===
+            ctx.update_task(progress=0.7, payload={"step": 3})
+
+            # Verify in-memory cache has LATEST update (third)
+            assert ctx._properties_cache.get("progress_percent") == 0.7
+            assert ctx._payload_cache.get("step") == 3
+
+            # Verify DB still has first update (both second and third deferred)
+            db.session.expire_all()
+            task_step2 = TaskDAO.find_one_or_none(uuid=task_uuid, skip_base_filter=True)
+            assert task_step2 is not None
+            assert task_step2.properties_dict.get("progress_percent") == 0.1
+            assert task_step2.payload_dict.get("step") == 1
+
+            # === Step 4: Wait for deferred timer to fire ===
+            time.sleep(throttle_interval + 0.5)
+
+            # Verify timer fired and wrote the LATEST update (third, not second)
+            db.session.expire_all()
+            task_step3 = TaskDAO.find_one_or_none(uuid=task_uuid, skip_base_filter=True)
+            assert task_step3 is not None
+            assert task_step3.properties_dict.get("progress_percent") == 0.7
+            assert task_step3.payload_dict.get("step") == 3
+
+        finally:
+            ctx._cancel_deferred_flush_timer()