amoghrajesh commented on code in PR #66586:
URL: https://github.com/apache/airflow/pull/66586#discussion_r3217655580


##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -47,15 +48,44 @@
 )
 from airflow.api_fastapi.core_api.security import GetUserDep
 from airflow.api_fastapi.core_api.services.public.common import BulkService
+from airflow.configuration import conf
 from airflow.listeners.listener import get_listener_manager
 from airflow.models.dag import DagModel
 from airflow.models.taskinstance import TaskInstance as TI
 from airflow.serialization.definitions.dag import SerializedDAG
+from airflow.state import get_state_backend
 from airflow.utils.state import TaskInstanceState
 
 log = structlog.get_logger(__name__)
 
 
+def _clear_task_state_on_success(tis: Sequence[TI], session: Session) -> None:
+    """Clear task state rows for each TI if clear_on_success is enabled."""
+    if not conf.getboolean("state_store", "clear_on_success", fallback=False):
+        return
+    backend = get_state_backend()
+    for ti in tis:
+        scope = TaskScope(
+            dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, 
map_index=ti.map_index or -1
+        )
+        try:
+            backend.clear(scope=scope, session=session)  # type: 
ignore[call-arg]  # @provide_session adds session kwarg at runtime; 
BaseStateBackend signature omits it so mypy can't see it

Review Comment:
   `session` only makes sense for `MetastoreStateBackend`-> it's a SQLA concern 
specific to this one implementation. Custom backends will not have a session at 
all. Exposing it on the base class or even in the override signature would be 
misleading: it implies all backends accept a session, which they do not. 
   
   The `# type: ignore[call-arg]` is for `get_state_backend()` which returns 
BaseStateBackend, which has no session parameter. I tried to make it 
non-keyword in `MetastoreStateBackend` and that doesn't help since mypy checks 
the call against the base class signature. The ignore is unavoidable here 
without adding session to BaseStateBackend, which would be wrong. Do you have 
any suggestions?



-- 
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]

Reply via email to