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]