from airflow.sdk import task
import random


@task
def my_task(**context):
    # Retrieve task_state_store from context
    task_state_store = context["task_state_store"]
    my_value = task_state_store.get("my_key", default="my_default_key")

    # Set the new value
    new_value = f"It is {random.randint(1, 12 + 1)} o'clock"
    task_state_store.set("my_key", new_value)

    # Delete the value
    task_state_store.delete("my_key")

    # Clear all store entries for the task
    task_state_store.clear()

Returns the stored JSON value, or the ``default`` value if the key does not exist.

.. code-block:: python

    value = task_state_store.get(
        "job_id", default="123456789"
    )  # returns the value associated with `job_id` or the default value

``set(key, value, *, retention=None)``

   # correct
   task_state_store.set("key", "val", retention=timedelta(days=7))

   # wrong — raises TypeError
   task_state_store.set("key", "val", retention=7)

from airflow.sdk import NEVER_EXPIRE

task_state_store.set("job_id", job_id, retention=NEVER_EXPIRE)

Deletes a single key. No-op if the key does not exist.

.. code-block:: python

    task_state_store.delete("job_id")

``clear(all_map_indices=False)``

# clear only this map index
task_state_store.clear()

# clear all map indices (fleet-wide)
task_state_store.clear(all_map_indices=True)

A common pattern for long-running external jobs: check whether a job ID is already stored before submitting, and use ``NEVER_EXPIRE`` so the key outlives
the default retention window.

.. code-block:: python

    from datetime import timedelta

    from airflow.sdk import DAG, task
    from airflow.sdk import NEVER_EXPIRE

    with DAG("spark_job_dag", schedule=None):

        @task
        def run_spark_job(**context):
            task_state_store = context["task_state_store"]

            # Check for an already-submitted job from a previous attempt.
            job_id = task_state_store.get("job_id")
            if job_id is None:
                job_id = spark_client.submit_job(...)
                # Store with NEVER_EXPIRE so the key is not garbage-collected before the job finishes
                task_state_store.set("job_id", job_id, retention=NEVER_EXPIRE)

            # Reattach to the job and wait for completion.
            result = spark_client.wait_for_completion(job_id)
            return result

On a retry, the task finds the stored ``job_id`` and reattaches instead of submitting a duplicate job. Another example of this sort of logic can be found in `example_task_state_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_task_state_store.py>`_.

For ``BaseOperator`` subclasses, the :class:`~airflow.sdk.bases.resumablemixin.ResumableJobMixin` encapsulates this pattern. It persists the external job ID to task state store after submission and, on retry, reconnects to an active job or resubmits if the prior job reached a terminal failure state.

Intra-task checkpointing

from airflow.sdk import DAG, task

with DAG("paginated_ingest", schedule="@daily"):

    @task
    def ingest_pages(**context):
        # Retrieve the task_state_store
        task_state_store = context["task_state_store"]
        raw = task_state_store.get("last_page")

        start_page = raw + 1 if raw is not None else 1

        for page in range(start_page, total_pages + 1):
            fetch_and_load(page)
            task_state_store.set("last_page", page)  # Update the task_state_store for reuse later

Task store can expose in-progress metrics for observability — row counts, status strings, or lightweight JSON payloads — without requiring XCom or an external system.

.. code-block:: python


    from airflow.sdk import DAG, task

    with DAG("row_ingest", schedule="@hourly"):

        @task
        def ingest_rows(**context):
            task_state_store = context["task_state_store"]
            total = 0

            for batch in get_batches():
                load(batch)
                total += len(batch)
                task_state_store.set(
                    "progress",
                    {"rows_loaded": total, "status": "running"},
                )

            task_state_store.set(
                "progress",
                {"rows_loaded": total, "status": "done"},
            )

The ``progress`` key is visible through the REST API and the Airflow UI while the task is running.


Sync vs. deferrable tasks
-------------------------

Task store behaves slightly differently depending on whether a task runs synchronously or uses the deferral mechanism.

Synchronous tasks

Once a task defers, the Triggerer handles continuity across poke cycles. Use task state store in deferrable tasks only when you need to survive an operator-initiated clear, not for normal poke continuity.


Mapped tasks
------------

When a task is dynamically mapped (``task.expand(...)``), each map index has its own task state store namespace. ``clear()`` without arguments clears the store only for the current index. ``clear(all_map_indices=True)`` wipes the store across every index of the task.

.. code-block:: python

    # Inside a mapped task — clear only this index
    task_state_store.clear()

    # Wipe store for all indices of this task
    task_state_store.clear(all_map_indices=True)


Automatic cleanup (``clear_on_success``)
----------------------------------------

When ``[state_store] clear_on_success = True``, all task state store keys for a task instance are automatically deleted when the task moves to the ``success`` state. This is useful for reducing storage when post-success observability is not needed.

.. note::

   ``clear_on_success`` clears **task state store only**. Asset store is scoped to the asset, not the task instance, and is never affected by this setting. Asset store persists across runs and must be cleared explicitly.

See :doc:`/administration-and-deployment/task-and-asset-state-store` for full configuration details.