amoghrajesh commented on code in PR #67299: URL: https://github.com/apache/airflow/pull/67299#discussion_r3356184022
########## airflow-core/docs/core-concepts/task-and-asset-store.rst: ########## @@ -0,0 +1,80 @@ + .. 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. + +.. _concepts:task-and-asset-store-overview: + +Task and Asset Store Overview +============================= + +.. versionadded:: 3.3 + +Airflow has always modeled tasks as stateless, idempotent units of work. A growing class of workloads, however, require a small amount of data to be persisted outside of a Task's return value, like a submitted job ID that must survive a worker crash, a watermark that advances run-by-run, or a row counter exposed for observability. Task store and Asset store fill that gap without touching the XCom or Variable systems. + +Task and Asset Store +-------------------- + +Airflow 3.3 ships two persistent key/value stores, differentiated by *what* they are scoped to: Review Comment: I don't think we need to mention versions, cos docs are versioned? ########## airflow-core/docs/core-concepts/task-and-asset-store.rst: ########## @@ -0,0 +1,80 @@ + .. 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. + +.. _concepts:task-and-asset-store-overview: + +Task and Asset Store Overview +============================= + +.. versionadded:: 3.3 + +Airflow has always modeled tasks as stateless, idempotent units of work. A growing class of workloads, however, require a small amount of data to be persisted outside of a Task's return value, like a submitted job ID that must survive a worker crash, a watermark that advances run-by-run, or a row counter exposed for observability. Task store and Asset store fill that gap without touching the XCom or Variable systems. + +Task and Asset Store +-------------------- + +Airflow 3.3 ships two persistent key/value stores, differentiated by *what* they are scoped to: + +.. list-table:: + :header-rows: 1 + :widths: 20 25 25 30 + + * - Store + - Scope + - Default lifetime + - Primary use case + * - **Task store** + - A single task Instance (dag_id + run_id + task_id + map_index) + - Configurable retention; cleared on task success when ``clear_on_success = True`` + - Survive retries, track in-flight jobs, checkpoint progress within a run, resume progress from checkpoint set by a past run + * - **Asset store** + - An asset (independent of any particular run) + - Persists indefinitely; removed only when the asset is deactivated + - Cross-run watermarks, incremental-load cursors, per-asset metadata + +Both stores accept string keys and JSON values. Values up to 64 KB are supported through the default metastore backend; larger payloads can be offloaded via a :ref:`custom worker-side backend <task-and-asset-store:worker-backends>`. + +When to use Task and Asset Store +-------------------------------- + +Use this table to choose the right mechanism for your use case. + +.. list-table:: + :header-rows: 1 + :widths: 22 78 + + * - Mechanism + - When to use it + * - **XCom** + - Pass data *between tasks* within a single Dag run (e.g. the output of one task consumed by a downstream task) or across different multiple Dag runs (referencing the data persisted from another run). XComs are cleared on retry, and should NOT be used to persist data across task retries or across runs. + * - **Variables** + - Dag-wide or installation-wide configuration that changes infrequently and is set by operators rather than by tasks themselves. Review Comment: Is it better to call this deployment wide? ########## airflow-core/docs/core-concepts/task-and-asset-store.rst: ########## @@ -0,0 +1,80 @@ + .. 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. + +.. _concepts:task-and-asset-store-overview: + +Task and Asset Store Overview +============================= + +.. versionadded:: 3.3 + +Airflow has always modeled tasks as stateless, idempotent units of work. A growing class of workloads, however, require a small amount of data to be persisted outside of a Task's return value, like a submitted job ID that must survive a worker crash, a watermark that advances run-by-run, or a row counter exposed for observability. Task store and Asset store fill that gap without touching the XCom or Variable systems. + +Task and Asset Store +-------------------- + +Airflow 3.3 ships two persistent key/value stores, differentiated by *what* they are scoped to: + +.. list-table:: + :header-rows: 1 + :widths: 20 25 25 30 + + * - Store + - Scope + - Default lifetime + - Primary use case + * - **Task store** + - A single task Instance (dag_id + run_id + task_id + map_index) + - Configurable retention; cleared on task success when ``clear_on_success = True`` + - Survive retries, track in-flight jobs, checkpoint progress within a run, resume progress from checkpoint set by a past run + * - **Asset store** + - An asset (independent of any particular run) + - Persists indefinitely; removed only when the asset is deactivated + - Cross-run watermarks, incremental-load cursors, per-asset metadata + +Both stores accept string keys and JSON values. Values up to 64 KB are supported through the default metastore backend; larger payloads can be offloaded via a :ref:`custom worker-side backend <task-and-asset-store:worker-backends>`. + +When to use Task and Asset Store +-------------------------------- + +Use this table to choose the right mechanism for your use case. + +.. list-table:: + :header-rows: 1 + :widths: 22 78 + + * - Mechanism + - When to use it + * - **XCom** + - Pass data *between tasks* within a single Dag run (e.g. the output of one task consumed by a downstream task) or across different multiple Dag runs (referencing the data persisted from another run). XComs are cleared on retry, and should NOT be used to persist data across task retries or across runs. + * - **Variables** + - Dag-wide or installation-wide configuration that changes infrequently and is set by operators rather than by tasks themselves. + * - **Task store** + - Data that must survive a worker crash or a retry within the **same run**. An external job ID written before a long-running job completes is a perfect use case for task store. Review Comment: ```suggestion - Data that must survive a worker crash or data that must survive across retries within the **same run**. An external job ID written before a long-running job completes is a perfect use case for task store. ``` ########## airflow-core/docs/core-concepts/task-store.rst: ########## @@ -0,0 +1,275 @@ + .. 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. + +.. _concepts:task-store: + +.. spelling:word-list:: + + intra + Intra + checkpointing + +Task Store +========== + +.. versionadded:: 3.3 + +Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. + +Data persisted via task store is accessed through the task context via ``context["task_store"]`` and exposes four methods: ``get``, ``set``, ``delete``, and ``clear``. + + +Accessing task store +-------------------- + +Inside any ``@task``-decorated function or ``BaseOperator.execute()`` method, task store is available through the ``context`` dictionary via the ``task_store`` key. From there, it can be used to retrieve, set, delete, or clear data for a specific key-value pair. In this example, the ``job_id`` is retrieved from task store, then updated, before being deleted. All data for that task is then removed using the ``clear`` method. + +.. code-block:: python + + from airflow.sdk import task + import random + + + @task + def my_task(**context): + # Retrieve task_store from context + task_store = context["task_store"] + my_value = task_store.get("my_key") + + # Set the new value + new_value = f"It is {random.randint(1, 12 + 1)} o'clock" + task_store.set("my_key", new_value) + + # Delete the value + task_store.delete("my_key") + + # Clear all store entries for the task + task_store.clear() + +Reference +------------- + +``get(key, default)`` +~~~~~~~~~~~~~~~~~~~~~~ + +Returns the stored string value, or the ``default`` value if the key does not exist. + +.. code-block:: python + + value = task_store.get("job_id", "123456789") # returns str or default value Review Comment: ```suggestion value = task_store.get("job_id", default="123456789") # returns the value associated with `job_id` or the default value ``` Unclear if we dont pass `default` here ########## airflow-core/docs/core-concepts/task-store.rst: ########## @@ -0,0 +1,275 @@ + .. 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. + +.. _concepts:task-store: + +.. spelling:word-list:: + + intra + Intra + checkpointing + +Task Store +========== + +.. versionadded:: 3.3 + +Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. + +Data persisted via task store is accessed through the task context via ``context["task_store"]`` and exposes four methods: ``get``, ``set``, ``delete``, and ``clear``. + + +Accessing task store +-------------------- + +Inside any ``@task``-decorated function or ``BaseOperator.execute()`` method, task store is available through the ``context`` dictionary via the ``task_store`` key. From there, it can be used to retrieve, set, delete, or clear data for a specific key-value pair. In this example, the ``job_id`` is retrieved from task store, then updated, before being deleted. All data for that task is then removed using the ``clear`` method. + +.. code-block:: python + + from airflow.sdk import task + import random + + + @task + def my_task(**context): + # Retrieve task_store from context + task_store = context["task_store"] + my_value = task_store.get("my_key") + + # Set the new value + new_value = f"It is {random.randint(1, 12 + 1)} o'clock" + task_store.set("my_key", new_value) + + # Delete the value + task_store.delete("my_key") + + # Clear all store entries for the task + task_store.clear() + +Reference +------------- + +``get(key, default)`` +~~~~~~~~~~~~~~~~~~~~~~ + +Returns the stored string value, or the ``default`` value if the key does not exist. + +.. code-block:: python + + value = task_store.get("job_id", "123456789") # returns str or default value + +``set(key, value, *, retention=None)`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Writes or overwrites a value for the specified key. Note, ``value`` can be any JSON-compatible type, except for ``None``. This includes: + +* ``str`` +* ``int`` +* ``float`` +* ``bool`` +* ``list`` +* ``dict`` + +The optional ``retention`` argument controls when the key expires: + +* ``timedelta(...)``: expire after the given duration from the time of the write (e.g. ``timedelta(hours=6)``). The expiry timestamp is computed on the worker in UTC before the value is sent to the API server. +* ``NEVER_EXPIRE``: the key never expires and is skipped by garbage collection, regardless of the global ``[state_store] default_retention_days`` setting. +* ``None`` (default): fall back to the global ``[state_store] default_retention_days`` config. + +.. important:: + + ``retention`` accepts only a :class:`~datetime.timedelta`, not a plain integer number of days. Passing an integer raises a ``TypeError``. + + .. code-block:: python + + # correct + task_store.set("key", "val", retention=timedelta(days=7)) + + # wrong — raises TypeError + task_store.set("key", "val", retention=7) + +``NEVER_EXPIRE`` sentinel +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Import ``NEVER_EXPIRE`` from ``airflow.sdk.execution_time.context``: + +.. code-block:: python + + from airflow.sdk import NEVER_EXPIRE + + task_store.set("job_id", job_id, retention=NEVER_EXPIRE) + +``delete(key)`` +~~~~~~~~~~~~~~~ + +Deletes a single key. No-op if the key does not exist. + +.. code-block:: python + + task_store.delete("job_id") + +``clear(all_map_indices=False)`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Deletes *all* task store keys for this task instance. + +For :doc:`mapped tasks </authoring-and-scheduling/dynamic-task-mapping>`, the default clears only the current map index. Pass ``all_map_indices=True`` to wipe the store across **every** mapped instance of the task (fleet-wide reset). + +.. code-block:: python + + # clear only this map index + task_store.clear() + + # clear all map indices (fleet-wide) + task_store.clear(all_map_indices=True) + + +Use Cases +--------- + +External job resumption +~~~~~~~~~~~~~~~~~~~~~~~ + +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_store = context["task_store"] + + # Check for an already-submitted job from a previous attempt. + job_id = task_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_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_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_task_store.py>`_. Review Comment: Should we mention the `ResumableJobMixin` class here? ########## airflow-core/docs/core-concepts/asset-store.rst: ########## @@ -0,0 +1,247 @@ + .. 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. + +.. _concepts:asset-store: + +.. spelling:word-list:: + + subscripted + subscripting + +Asset Store +=========== + +.. versionadded:: 3.3 + +Asset store is a persistent key/value store scoped to an *asset*, independent of any particular DAG run. Unlike :doc:`task store </core-concepts/task-store>`, which is tied to a single task instance, asset store persists across runs and is logically owned by the asset itself. It is the natural home for cross-run metadata such as watermarks, incremental-load cursors, and per-asset configuration. + +Asset store is accessed through the task context via ``context["asset_store"]``. + + +When is ``asset_store`` available? +------------------------------------ + +When using asset store within a task, ``context["asset_store"]`` is populated for **concrete** :class:`~airflow.sdk.definitions.asset.Asset` inlets and outlets. A task must declare at least one concrete inlet or outlet for ``asset_store`` to contain any entries. + +If using asset store in a ``BaseEventTrigger``, the ``self.asset_store`` parameter can be used within the ``BaseEventTrigger``. It can be subscripted in the same way that ``context["asset_store"]`` can be. + +.. warning:: + + **Outlets-only tasks**: if a task declares only ``outlets`` (no ``inlets``), ``context["asset_store"][my_asset]`` may raise a ``KeyError`` at runtime. The workaround is to declare the asset in **both** ``inlets`` and ``outlets``. + + .. code-block:: python + + # my_asset defined above ... + + + @task(inlets=[my_asset], outlets=[my_asset]) + def write_asset(**context): + context["asset_store"][my_asset].set("watermark", "2024-01-01") + + This known issue will be resolved in a future release. + + +Accessing asset store using ``context`` +--------------------------------------- + +An asset can be brought into "scope" (for lack of a better phrase) by including it in ``inlets`` (or both ``inlets`` and ``outlets``). Then subscript ``context["asset_store"]`` with the asset object to retrieve the asset store. Review Comment: ```suggestion An asset becomes available through context["asset_store"] when it is included in inlets (or in both inlets and outlets). You can then retrieve its asset store by subscripting context["asset_store"] with the asset object. ``` ########## airflow-core/docs/core-concepts/asset-store.rst: ########## @@ -0,0 +1,247 @@ + .. 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. + +.. _concepts:asset-store: + +.. spelling:word-list:: + + subscripted + subscripting + +Asset Store +=========== + +.. versionadded:: 3.3 + +Asset store is a persistent key/value store scoped to an *asset*, independent of any particular DAG run. Unlike :doc:`task store </core-concepts/task-store>`, which is tied to a single task instance, asset store persists across runs and is logically owned by the asset itself. It is the natural home for cross-run metadata such as watermarks, incremental-load cursors, and per-asset configuration. + +Asset store is accessed through the task context via ``context["asset_store"]``. + + +When is ``asset_store`` available? +------------------------------------ + +When using asset store within a task, ``context["asset_store"]`` is populated for **concrete** :class:`~airflow.sdk.definitions.asset.Asset` inlets and outlets. A task must declare at least one concrete inlet or outlet for ``asset_store`` to contain any entries. + +If using asset store in a ``BaseEventTrigger``, the ``self.asset_store`` parameter can be used within the ``BaseEventTrigger``. It can be subscripted in the same way that ``context["asset_store"]`` can be. + +.. warning:: + + **Outlets-only tasks**: if a task declares only ``outlets`` (no ``inlets``), ``context["asset_store"][my_asset]`` may raise a ``KeyError`` at runtime. The workaround is to declare the asset in **both** ``inlets`` and ``outlets``. + + .. code-block:: python + + # my_asset defined above ... + + + @task(inlets=[my_asset], outlets=[my_asset]) + def write_asset(**context): + context["asset_store"][my_asset].set("watermark", "2024-01-01") + + This known issue will be resolved in a future release. + + +Accessing asset store using ``context`` +--------------------------------------- + +An asset can be brought into "scope" (for lack of a better phrase) by including it in ``inlets`` (or both ``inlets`` and ``outlets``). Then subscript ``context["asset_store"]`` with the asset object to retrieve the asset store. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + my_asset = Asset("my_data", uri="s3://bucket/my_data") + + with DAG("example_asset_store", schedule=None): + + @task(inlets=[my_asset], outlets=[my_asset]) + def process(**context): + asset_store = context["asset_store"][my_asset] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +To see asset store in-action in a real DAG, checkout the DAG in `example_asset_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_asset_store.py>`_. + +Accessing asset store in a ``BaseEventTrigger`` +----------------------------------------------- + +When building Triggers used for asset "watching", asset store can be retrieved using the ``self.asset_store`` attribute. + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(name="my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +In the example above, ``my_data`` is created using the ``name`` However, the ``uri`` can also be used: + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(uri="s3://bucket/my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +Single-inlet shorthand +~~~~~~~~~~~~~~~~~~~~~~~ + +For tasks with exactly **one** concrete inlet, you can call ``get``, ``set``, ``delete``, and ``clear`` directly on ``context["asset_store"]`` without subscripting. + +.. code-block:: python + + @task(inlets=[my_asset], outlets=[my_asset]) + def process_single(**context): + asset_store = context["asset_store"] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +If the task has more than one concrete inlet, calling the shorthand raises a ``ValueError``. Use the subscript form (``context["asset_store"][my_asset]``) whenever a task has multiple inlets. + + +API reference +------------- + +The following methods are available on both the per-asset accessor (``context["asset_store"][my_asset]``), the shorthand (``context["asset_store"]``) when the task has exactly one inlet, and when using the ``self.asset_store`` attribute. + +``get(key)`` +~~~~~~~~~~~~ + +Returns the stored JSON value, or ``None`` if the key does not exist. + +.. code-block:: python + + # Using context + watermark = context["asset_store"][my_asset].get("watermark") Review Comment: ```suggestion # Using context watermark = context["asset_store"][my_asset].get("watermark", default="initial_watermark") ``` ########## airflow-core/docs/core-concepts/asset-store.rst: ########## @@ -0,0 +1,247 @@ + .. 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. + +.. _concepts:asset-store: + +.. spelling:word-list:: + + subscripted + subscripting + +Asset Store +=========== + +.. versionadded:: 3.3 + +Asset store is a persistent key/value store scoped to an *asset*, independent of any particular DAG run. Unlike :doc:`task store </core-concepts/task-store>`, which is tied to a single task instance, asset store persists across runs and is logically owned by the asset itself. It is the natural home for cross-run metadata such as watermarks, incremental-load cursors, and per-asset configuration. + +Asset store is accessed through the task context via ``context["asset_store"]``. + + +When is ``asset_store`` available? +------------------------------------ + +When using asset store within a task, ``context["asset_store"]`` is populated for **concrete** :class:`~airflow.sdk.definitions.asset.Asset` inlets and outlets. A task must declare at least one concrete inlet or outlet for ``asset_store`` to contain any entries. + +If using asset store in a ``BaseEventTrigger``, the ``self.asset_store`` parameter can be used within the ``BaseEventTrigger``. It can be subscripted in the same way that ``context["asset_store"]`` can be. + +.. warning:: + + **Outlets-only tasks**: if a task declares only ``outlets`` (no ``inlets``), ``context["asset_store"][my_asset]`` may raise a ``KeyError`` at runtime. The workaround is to declare the asset in **both** ``inlets`` and ``outlets``. + + .. code-block:: python + + # my_asset defined above ... + + + @task(inlets=[my_asset], outlets=[my_asset]) + def write_asset(**context): + context["asset_store"][my_asset].set("watermark", "2024-01-01") + + This known issue will be resolved in a future release. + + +Accessing asset store using ``context`` +--------------------------------------- + +An asset can be brought into "scope" (for lack of a better phrase) by including it in ``inlets`` (or both ``inlets`` and ``outlets``). Then subscript ``context["asset_store"]`` with the asset object to retrieve the asset store. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + my_asset = Asset("my_data", uri="s3://bucket/my_data") + + with DAG("example_asset_store", schedule=None): + + @task(inlets=[my_asset], outlets=[my_asset]) + def process(**context): + asset_store = context["asset_store"][my_asset] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +To see asset store in-action in a real DAG, checkout the DAG in `example_asset_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_asset_store.py>`_. + +Accessing asset store in a ``BaseEventTrigger`` +----------------------------------------------- + +When building Triggers used for asset "watching", asset store can be retrieved using the ``self.asset_store`` attribute. + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(name="my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +In the example above, ``my_data`` is created using the ``name`` However, the ``uri`` can also be used: + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(uri="s3://bucket/my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +Single-inlet shorthand +~~~~~~~~~~~~~~~~~~~~~~~ + +For tasks with exactly **one** concrete inlet, you can call ``get``, ``set``, ``delete``, and ``clear`` directly on ``context["asset_store"]`` without subscripting. + +.. code-block:: python + + @task(inlets=[my_asset], outlets=[my_asset]) + def process_single(**context): + asset_store = context["asset_store"] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +If the task has more than one concrete inlet, calling the shorthand raises a ``ValueError``. Use the subscript form (``context["asset_store"][my_asset]``) whenever a task has multiple inlets. + + +API reference +------------- + +The following methods are available on both the per-asset accessor (``context["asset_store"][my_asset]``), the shorthand (``context["asset_store"]``) when the task has exactly one inlet, and when using the ``self.asset_store`` attribute. + +``get(key)`` Review Comment: ```suggestion ``get(key, default)`` ``` ########## airflow-core/docs/core-concepts/task-and-asset-store.rst: ########## @@ -0,0 +1,80 @@ + .. 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. + +.. _concepts:task-and-asset-store-overview: + +Task and Asset Store Overview +============================= + +.. versionadded:: 3.3 + +Airflow has always modeled tasks as stateless, idempotent units of work. A growing class of workloads, however, require a small amount of data to be persisted outside of a Task's return value, like a submitted job ID that must survive a worker crash, a watermark that advances run-by-run, or a row counter exposed for observability. Task store and Asset store fill that gap without touching the XCom or Variable systems. + +Task and Asset Store +-------------------- + +Airflow 3.3 ships two persistent key/value stores, differentiated by *what* they are scoped to: + +.. list-table:: + :header-rows: 1 + :widths: 20 25 25 30 + + * - Store + - Scope + - Default lifetime + - Primary use case + * - **Task store** + - A single task Instance (dag_id + run_id + task_id + map_index) + - Configurable retention; cleared on task success when ``clear_on_success = True`` + - Survive retries, track in-flight jobs, checkpoint progress within a run, resume progress from checkpoint set by a past run + * - **Asset store** + - An asset (independent of any particular run) + - Persists indefinitely; removed only when the asset is deactivated + - Cross-run watermarks, incremental-load cursors, per-asset metadata + +Both stores accept string keys and JSON values. Values up to 64 KB are supported through the default metastore backend; larger payloads can be offloaded via a :ref:`custom worker-side backend <task-and-asset-store:worker-backends>`. Review Comment: ```suggestion Both stores accept JSON-able values. Values up to 64 KB are supported through the default metastore backend; larger payloads can be offloaded via a :ref:`custom worker-side backend <task-and-asset-store:worker-backends>`. ``` ########## airflow-core/docs/core-concepts/task-and-asset-store.rst: ########## @@ -0,0 +1,80 @@ + .. 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. + +.. _concepts:task-and-asset-store-overview: + +Task and Asset Store Overview +============================= + +.. versionadded:: 3.3 + +Airflow has always modeled tasks as stateless, idempotent units of work. A growing class of workloads, however, require a small amount of data to be persisted outside of a Task's return value, like a submitted job ID that must survive a worker crash, a watermark that advances run-by-run, or a row counter exposed for observability. Task store and Asset store fill that gap without touching the XCom or Variable systems. Review Comment: ```suggestion Airflow has always modeled tasks as stateless, idempotent units of work. A growing class of workloads, however, require some amount of data to be persisted outside of a yask's return value, like a submitted job ID that must survive a worker crash, a watermark that advances run-by-run, or a row counter exposed for observability. Task store and Asset store fill that gap without touching the XCom or Variable systems. ``` ########## airflow-core/docs/core-concepts/task-store.rst: ########## @@ -0,0 +1,275 @@ + .. 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. + +.. _concepts:task-store: + +.. spelling:word-list:: + + intra + Intra + checkpointing + +Task Store +========== + +.. versionadded:: 3.3 + +Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. + +Data persisted via task store is accessed through the task context via ``context["task_store"]`` and exposes four methods: ``get``, ``set``, ``delete``, and ``clear``. + + +Accessing task store +-------------------- + +Inside any ``@task``-decorated function or ``BaseOperator.execute()`` method, task store is available through the ``context`` dictionary via the ``task_store`` key. From there, it can be used to retrieve, set, delete, or clear data for a specific key-value pair. In this example, the ``job_id`` is retrieved from task store, then updated, before being deleted. All data for that task is then removed using the ``clear`` method. + +.. code-block:: python + + from airflow.sdk import task + import random + + + @task + def my_task(**context): + # Retrieve task_store from context + task_store = context["task_store"] + my_value = task_store.get("my_key") Review Comment: ```suggestion my_value = task_store.get("my_key", default="my_default_key") ``` ########## airflow-core/docs/core-concepts/task-store.rst: ########## @@ -0,0 +1,275 @@ + .. 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. + +.. _concepts:task-store: + +.. spelling:word-list:: + + intra + Intra + checkpointing + +Task Store +========== + +.. versionadded:: 3.3 + +Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. + +Data persisted via task store is accessed through the task context via ``context["task_store"]`` and exposes four methods: ``get``, ``set``, ``delete``, and ``clear``. + + +Accessing task store +-------------------- + +Inside any ``@task``-decorated function or ``BaseOperator.execute()`` method, task store is available through the ``context`` dictionary via the ``task_store`` key. From there, it can be used to retrieve, set, delete, or clear data for a specific key-value pair. In this example, the ``job_id`` is retrieved from task store, then updated, before being deleted. All data for that task is then removed using the ``clear`` method. + +.. code-block:: python + + from airflow.sdk import task + import random + + + @task + def my_task(**context): + # Retrieve task_store from context + task_store = context["task_store"] + my_value = task_store.get("my_key") + + # Set the new value + new_value = f"It is {random.randint(1, 12 + 1)} o'clock" + task_store.set("my_key", new_value) + + # Delete the value + task_store.delete("my_key") + + # Clear all store entries for the task + task_store.clear() + +Reference +------------- + +``get(key, default)`` +~~~~~~~~~~~~~~~~~~~~~~ + +Returns the stored string value, or the ``default`` value if the key does not exist. + +.. code-block:: python + + value = task_store.get("job_id", "123456789") # returns str or default value + +``set(key, value, *, retention=None)`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Writes or overwrites a value for the specified key. Note, ``value`` can be any JSON-compatible type, except for ``None``. This includes: + +* ``str`` +* ``int`` +* ``float`` +* ``bool`` +* ``list`` +* ``dict`` + +The optional ``retention`` argument controls when the key expires: + +* ``timedelta(...)``: expire after the given duration from the time of the write (e.g. ``timedelta(hours=6)``). The expiry timestamp is computed on the worker in UTC before the value is sent to the API server. Review Comment: UTC is irrelevant detail? ########## airflow-core/docs/core-concepts/task-store.rst: ########## @@ -0,0 +1,275 @@ + .. 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. + +.. _concepts:task-store: + +.. spelling:word-list:: + + intra + Intra + checkpointing + +Task Store +========== + +.. versionadded:: 3.3 + +Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. + +Data persisted via task store is accessed through the task context via ``context["task_store"]`` and exposes four methods: ``get``, ``set``, ``delete``, and ``clear``. + + +Accessing task store +-------------------- + +Inside any ``@task``-decorated function or ``BaseOperator.execute()`` method, task store is available through the ``context`` dictionary via the ``task_store`` key. From there, it can be used to retrieve, set, delete, or clear data for a specific key-value pair. In this example, the ``job_id`` is retrieved from task store, then updated, before being deleted. All data for that task is then removed using the ``clear`` method. + +.. code-block:: python + + from airflow.sdk import task + import random + + + @task + def my_task(**context): + # Retrieve task_store from context + task_store = context["task_store"] + my_value = task_store.get("my_key") + + # Set the new value + new_value = f"It is {random.randint(1, 12 + 1)} o'clock" + task_store.set("my_key", new_value) + + # Delete the value + task_store.delete("my_key") + + # Clear all store entries for the task + task_store.clear() + +Reference +------------- + +``get(key, default)`` +~~~~~~~~~~~~~~~~~~~~~~ + +Returns the stored string value, or the ``default`` value if the key does not exist. + +.. code-block:: python + + value = task_store.get("job_id", "123456789") # returns str or default value + +``set(key, value, *, retention=None)`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Writes or overwrites a value for the specified key. Note, ``value`` can be any JSON-compatible type, except for ``None``. This includes: + +* ``str`` +* ``int`` +* ``float`` +* ``bool`` +* ``list`` +* ``dict`` + +The optional ``retention`` argument controls when the key expires: + +* ``timedelta(...)``: expire after the given duration from the time of the write (e.g. ``timedelta(hours=6)``). The expiry timestamp is computed on the worker in UTC before the value is sent to the API server. +* ``NEVER_EXPIRE``: the key never expires and is skipped by garbage collection, regardless of the global ``[state_store] default_retention_days`` setting. Review Comment: ```suggestion * ``NEVER_EXPIRE``: the key never expires and is skipped during garbage collection, regardless of the global ``[state_store] default_retention_days`` setting. ``` ########## airflow-core/docs/core-concepts/task-store.rst: ########## @@ -0,0 +1,275 @@ + .. 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. + +.. _concepts:task-store: + +.. spelling:word-list:: + + intra + Intra + checkpointing + +Task Store +========== + +.. versionadded:: 3.3 + +Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. + +Data persisted via task store is accessed through the task context via ``context["task_store"]`` and exposes four methods: ``get``, ``set``, ``delete``, and ``clear``. + + +Accessing task store +-------------------- + +Inside any ``@task``-decorated function or ``BaseOperator.execute()`` method, task store is available through the ``context`` dictionary via the ``task_store`` key. From there, it can be used to retrieve, set, delete, or clear data for a specific key-value pair. In this example, the ``job_id`` is retrieved from task store, then updated, before being deleted. All data for that task is then removed using the ``clear`` method. + +.. code-block:: python + + from airflow.sdk import task + import random + + + @task + def my_task(**context): + # Retrieve task_store from context + task_store = context["task_store"] + my_value = task_store.get("my_key") + + # Set the new value + new_value = f"It is {random.randint(1, 12 + 1)} o'clock" + task_store.set("my_key", new_value) + + # Delete the value + task_store.delete("my_key") + + # Clear all store entries for the task + task_store.clear() + +Reference +------------- + +``get(key, default)`` +~~~~~~~~~~~~~~~~~~~~~~ + +Returns the stored string value, or the ``default`` value if the key does not exist. + +.. code-block:: python + + value = task_store.get("job_id", "123456789") # returns str or default value + +``set(key, value, *, retention=None)`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Writes or overwrites a value for the specified key. Note, ``value`` can be any JSON-compatible type, except for ``None``. This includes: + +* ``str`` +* ``int`` +* ``float`` +* ``bool`` +* ``list`` +* ``dict`` + +The optional ``retention`` argument controls when the key expires: + +* ``timedelta(...)``: expire after the given duration from the time of the write (e.g. ``timedelta(hours=6)``). The expiry timestamp is computed on the worker in UTC before the value is sent to the API server. +* ``NEVER_EXPIRE``: the key never expires and is skipped by garbage collection, regardless of the global ``[state_store] default_retention_days`` setting. +* ``None`` (default): fall back to the global ``[state_store] default_retention_days`` config. + +.. important:: + + ``retention`` accepts only a :class:`~datetime.timedelta`, not a plain integer number of days. Passing an integer raises a ``TypeError``. + + .. code-block:: python + + # correct + task_store.set("key", "val", retention=timedelta(days=7)) + + # wrong — raises TypeError + task_store.set("key", "val", retention=7) + +``NEVER_EXPIRE`` sentinel +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Import ``NEVER_EXPIRE`` from ``airflow.sdk.execution_time.context``: + +.. code-block:: python + + from airflow.sdk import NEVER_EXPIRE + + task_store.set("job_id", job_id, retention=NEVER_EXPIRE) + +``delete(key)`` +~~~~~~~~~~~~~~~ + +Deletes a single key. No-op if the key does not exist. + +.. code-block:: python + + task_store.delete("job_id") + +``clear(all_map_indices=False)`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Deletes *all* task store keys for this task instance. + +For :doc:`mapped tasks </authoring-and-scheduling/dynamic-task-mapping>`, the default clears only the current map index. Pass ``all_map_indices=True`` to wipe the store across **every** mapped instance of the task (fleet-wide reset). + +.. code-block:: python + + # clear only this map index + task_store.clear() + + # clear all map indices (fleet-wide) + task_store.clear(all_map_indices=True) + + +Use Cases Review Comment: ```suggestion Some Example Use Cases ``` IDK, is this better? ########## airflow-core/docs/core-concepts/task-store.rst: ########## @@ -0,0 +1,275 @@ + .. 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. + +.. _concepts:task-store: + +.. spelling:word-list:: + + intra + Intra + checkpointing + +Task Store +========== + +.. versionadded:: 3.3 + +Task store is a persistent key/value store scoped to a single task instance (``dag_id`` + ``run_id`` + ``task_id`` + ``map_index``). It survives worker crashes and task retries within the same Dag run, making it suitable for storing external job IDs, intra-task checkpoints, and progress metadata. + +Data persisted via task store is accessed through the task context via ``context["task_store"]`` and exposes four methods: ``get``, ``set``, ``delete``, and ``clear``. + + +Accessing task store +-------------------- + +Inside any ``@task``-decorated function or ``BaseOperator.execute()`` method, task store is available through the ``context`` dictionary via the ``task_store`` key. From there, it can be used to retrieve, set, delete, or clear data for a specific key-value pair. In this example, the ``job_id`` is retrieved from task store, then updated, before being deleted. All data for that task is then removed using the ``clear`` method. + +.. code-block:: python + + from airflow.sdk import task + import random + + + @task + def my_task(**context): + # Retrieve task_store from context + task_store = context["task_store"] + my_value = task_store.get("my_key") + + # Set the new value + new_value = f"It is {random.randint(1, 12 + 1)} o'clock" + task_store.set("my_key", new_value) + + # Delete the value + task_store.delete("my_key") + + # Clear all store entries for the task + task_store.clear() + +Reference +------------- + +``get(key, default)`` +~~~~~~~~~~~~~~~~~~~~~~ + +Returns the stored string value, or the ``default`` value if the key does not exist. + +.. code-block:: python + + value = task_store.get("job_id", "123456789") # returns str or default value + +``set(key, value, *, retention=None)`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Writes or overwrites a value for the specified key. Note, ``value`` can be any JSON-compatible type, except for ``None``. This includes: + +* ``str`` +* ``int`` +* ``float`` +* ``bool`` +* ``list`` +* ``dict`` + +The optional ``retention`` argument controls when the key expires: + +* ``timedelta(...)``: expire after the given duration from the time of the write (e.g. ``timedelta(hours=6)``). The expiry timestamp is computed on the worker in UTC before the value is sent to the API server. +* ``NEVER_EXPIRE``: the key never expires and is skipped by garbage collection, regardless of the global ``[state_store] default_retention_days`` setting. +* ``None`` (default): fall back to the global ``[state_store] default_retention_days`` config. + +.. important:: + + ``retention`` accepts only a :class:`~datetime.timedelta`, not a plain integer number of days. Passing an integer raises a ``TypeError``. + + .. code-block:: python + + # correct + task_store.set("key", "val", retention=timedelta(days=7)) + + # wrong — raises TypeError + task_store.set("key", "val", retention=7) + +``NEVER_EXPIRE`` sentinel +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Import ``NEVER_EXPIRE`` from ``airflow.sdk.execution_time.context``: + +.. code-block:: python + + from airflow.sdk import NEVER_EXPIRE + + task_store.set("job_id", job_id, retention=NEVER_EXPIRE) + +``delete(key)`` +~~~~~~~~~~~~~~~ + +Deletes a single key. No-op if the key does not exist. + +.. code-block:: python + + task_store.delete("job_id") + +``clear(all_map_indices=False)`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Deletes *all* task store keys for this task instance. + +For :doc:`mapped tasks </authoring-and-scheduling/dynamic-task-mapping>`, the default clears only the current map index. Pass ``all_map_indices=True`` to wipe the store across **every** mapped instance of the task (fleet-wide reset). + +.. code-block:: python + + # clear only this map index + task_store.clear() + + # clear all map indices (fleet-wide) + task_store.clear(all_map_indices=True) + + +Use Cases +--------- + +External job resumption +~~~~~~~~~~~~~~~~~~~~~~~ + +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_store = context["task_store"] + + # Check for an already-submitted job from a previous attempt. + job_id = task_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_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_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_task_store.py>`_. + +Intra-task checkpointing +~~~~~~~~~~~~~~~~~~~~~~~~ + +For tasks that process paginated or batched data, store the last-completed offset so a retry can resume mid-stream rather than restarting from the beginning. + +.. code-block:: python + + from airflow.sdk import DAG, task + + with DAG("paginated_ingest", schedule="@daily"): + + @task + def ingest_pages(**context): + # Retrieve the task_store + task_store = context["task_store"] + raw = task_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_store.set("last_page", page) # Update the task_store for reuse later + + +On a retry, the task reads ``last_page`` and skips pages that were already processed. + +Progress metadata +~~~~~~~~~~~~~~~~~ + +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 + + import json Review Comment: ```suggestion ``` ########## airflow-core/docs/core-concepts/asset-store.rst: ########## @@ -0,0 +1,247 @@ + .. 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. + +.. _concepts:asset-store: + +.. spelling:word-list:: + + subscripted + subscripting + +Asset Store +=========== + +.. versionadded:: 3.3 + +Asset store is a persistent key/value store scoped to an *asset*, independent of any particular DAG run. Unlike :doc:`task store </core-concepts/task-store>`, which is tied to a single task instance, asset store persists across runs and is logically owned by the asset itself. It is the natural home for cross-run metadata such as watermarks, incremental-load cursors, and per-asset configuration. + +Asset store is accessed through the task context via ``context["asset_store"]``. + + +When is ``asset_store`` available? +------------------------------------ + +When using asset store within a task, ``context["asset_store"]`` is populated for **concrete** :class:`~airflow.sdk.definitions.asset.Asset` inlets and outlets. A task must declare at least one concrete inlet or outlet for ``asset_store`` to contain any entries. + +If using asset store in a ``BaseEventTrigger``, the ``self.asset_store`` parameter can be used within the ``BaseEventTrigger``. It can be subscripted in the same way that ``context["asset_store"]`` can be. + +.. warning:: + + **Outlets-only tasks**: if a task declares only ``outlets`` (no ``inlets``), ``context["asset_store"][my_asset]`` may raise a ``KeyError`` at runtime. The workaround is to declare the asset in **both** ``inlets`` and ``outlets``. + + .. code-block:: python + + # my_asset defined above ... + + + @task(inlets=[my_asset], outlets=[my_asset]) + def write_asset(**context): + context["asset_store"][my_asset].set("watermark", "2024-01-01") + + This known issue will be resolved in a future release. + + +Accessing asset store using ``context`` +--------------------------------------- + +An asset can be brought into "scope" (for lack of a better phrase) by including it in ``inlets`` (or both ``inlets`` and ``outlets``). Then subscript ``context["asset_store"]`` with the asset object to retrieve the asset store. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + my_asset = Asset("my_data", uri="s3://bucket/my_data") + + with DAG("example_asset_store", schedule=None): + + @task(inlets=[my_asset], outlets=[my_asset]) + def process(**context): + asset_store = context["asset_store"][my_asset] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +To see asset store in-action in a real DAG, checkout the DAG in `example_asset_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_asset_store.py>`_. + +Accessing asset store in a ``BaseEventTrigger`` +----------------------------------------------- + +When building Triggers used for asset "watching", asset store can be retrieved using the ``self.asset_store`` attribute. + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(name="my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +In the example above, ``my_data`` is created using the ``name`` However, the ``uri`` can also be used: + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(uri="s3://bucket/my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") Review Comment: I suggest you move documentation related to `BaseEventTrigger` in a follow up OR in the PR where you add support for that since that PR has not been merged yet, we can close this initial framework sooner that way. ########## airflow-core/docs/core-concepts/asset-store.rst: ########## @@ -0,0 +1,247 @@ + .. 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. + +.. _concepts:asset-store: + +.. spelling:word-list:: + + subscripted + subscripting + +Asset Store +=========== + +.. versionadded:: 3.3 + +Asset store is a persistent key/value store scoped to an *asset*, independent of any particular DAG run. Unlike :doc:`task store </core-concepts/task-store>`, which is tied to a single task instance, asset store persists across runs and is logically owned by the asset itself. It is the natural home for cross-run metadata such as watermarks, incremental-load cursors, and per-asset configuration. + +Asset store is accessed through the task context via ``context["asset_store"]``. + + +When is ``asset_store`` available? +------------------------------------ + +When using asset store within a task, ``context["asset_store"]`` is populated for **concrete** :class:`~airflow.sdk.definitions.asset.Asset` inlets and outlets. A task must declare at least one concrete inlet or outlet for ``asset_store`` to contain any entries. + +If using asset store in a ``BaseEventTrigger``, the ``self.asset_store`` parameter can be used within the ``BaseEventTrigger``. It can be subscripted in the same way that ``context["asset_store"]`` can be. + +.. warning:: + + **Outlets-only tasks**: if a task declares only ``outlets`` (no ``inlets``), ``context["asset_store"][my_asset]`` may raise a ``KeyError`` at runtime. The workaround is to declare the asset in **both** ``inlets`` and ``outlets``. + + .. code-block:: python + + # my_asset defined above ... + + + @task(inlets=[my_asset], outlets=[my_asset]) + def write_asset(**context): + context["asset_store"][my_asset].set("watermark", "2024-01-01") + + This known issue will be resolved in a future release. + + +Accessing asset store using ``context`` +--------------------------------------- + +An asset can be brought into "scope" (for lack of a better phrase) by including it in ``inlets`` (or both ``inlets`` and ``outlets``). Then subscript ``context["asset_store"]`` with the asset object to retrieve the asset store. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + my_asset = Asset("my_data", uri="s3://bucket/my_data") + + with DAG("example_asset_store", schedule=None): + + @task(inlets=[my_asset], outlets=[my_asset]) + def process(**context): + asset_store = context["asset_store"][my_asset] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +To see asset store in-action in a real DAG, checkout the DAG in `example_asset_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_asset_store.py>`_. + +Accessing asset store in a ``BaseEventTrigger`` +----------------------------------------------- + +When building Triggers used for asset "watching", asset store can be retrieved using the ``self.asset_store`` attribute. + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(name="my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +In the example above, ``my_data`` is created using the ``name`` However, the ``uri`` can also be used: + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(uri="s3://bucket/my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +Single-inlet shorthand +~~~~~~~~~~~~~~~~~~~~~~~ + +For tasks with exactly **one** concrete inlet, you can call ``get``, ``set``, ``delete``, and ``clear`` directly on ``context["asset_store"]`` without subscripting. + +.. code-block:: python + + @task(inlets=[my_asset], outlets=[my_asset]) + def process_single(**context): + asset_store = context["asset_store"] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +If the task has more than one concrete inlet, calling the shorthand raises a ``ValueError``. Use the subscript form (``context["asset_store"][my_asset]``) whenever a task has multiple inlets. + + +API reference +------------- + +The following methods are available on both the per-asset accessor (``context["asset_store"][my_asset]``), the shorthand (``context["asset_store"]``) when the task has exactly one inlet, and when using the ``self.asset_store`` attribute. + +``get(key)`` +~~~~~~~~~~~~ + +Returns the stored JSON value, or ``None`` if the key does not exist. + +.. code-block:: python + + # Using context + watermark = context["asset_store"][my_asset].get("watermark") + + # Using self.asset_store + watermark = self.asset_store[my_asset].get("watermark") + +``set(key, value)`` +~~~~~~~~~~~~~~~~~~~ + +Writes or overwrites a key-value pair. Unlike task store, asset store has no ``retention`` parameter. Values persist until explicitly deleted or until the asset is deactivated. Like with task store, ``value`` can be any JSON-compatible type, except for ``None``. This includes: + +* ``str`` +* ``int`` +* ``float`` +* ``bool`` +* ``list`` +* ``dict`` + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].set("watermark", "2024-06-01T00:00:00Z") + + # Using self.asset_store + self.asset_store[my_asset].set("watermark", "2024-06-01T00:00:00Z") + +``delete(key)`` +~~~~~~~~~~~~~~~ + +Deletes a single key. No-op if the key does not exist. + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].delete("watermark") + + # Using self.asset_store + self.asset_store[my_asset].delete("watermark") + +``clear()`` +~~~~~~~~~~~ + +Deletes *all* asset store keys for the asset. + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].clear() + + # Using self.asset_store + self.asset_store[my_asset].clear() + +Use cases Review Comment: ```suggestion Some Example Use cases ``` ########## airflow-core/docs/core-concepts/asset-store.rst: ########## @@ -0,0 +1,247 @@ + .. 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. + +.. _concepts:asset-store: + +.. spelling:word-list:: + + subscripted + subscripting + +Asset Store +=========== + +.. versionadded:: 3.3 + +Asset store is a persistent key/value store scoped to an *asset*, independent of any particular DAG run. Unlike :doc:`task store </core-concepts/task-store>`, which is tied to a single task instance, asset store persists across runs and is logically owned by the asset itself. It is the natural home for cross-run metadata such as watermarks, incremental-load cursors, and per-asset configuration. + +Asset store is accessed through the task context via ``context["asset_store"]``. + + +When is ``asset_store`` available? +------------------------------------ + +When using asset store within a task, ``context["asset_store"]`` is populated for **concrete** :class:`~airflow.sdk.definitions.asset.Asset` inlets and outlets. A task must declare at least one concrete inlet or outlet for ``asset_store`` to contain any entries. + +If using asset store in a ``BaseEventTrigger``, the ``self.asset_store`` parameter can be used within the ``BaseEventTrigger``. It can be subscripted in the same way that ``context["asset_store"]`` can be. + +.. warning:: + + **Outlets-only tasks**: if a task declares only ``outlets`` (no ``inlets``), ``context["asset_store"][my_asset]`` may raise a ``KeyError`` at runtime. The workaround is to declare the asset in **both** ``inlets`` and ``outlets``. + + .. code-block:: python + + # my_asset defined above ... + + + @task(inlets=[my_asset], outlets=[my_asset]) + def write_asset(**context): + context["asset_store"][my_asset].set("watermark", "2024-01-01") + + This known issue will be resolved in a future release. + + +Accessing asset store using ``context`` +--------------------------------------- + +An asset can be brought into "scope" (for lack of a better phrase) by including it in ``inlets`` (or both ``inlets`` and ``outlets``). Then subscript ``context["asset_store"]`` with the asset object to retrieve the asset store. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + my_asset = Asset("my_data", uri="s3://bucket/my_data") + + with DAG("example_asset_store", schedule=None): + + @task(inlets=[my_asset], outlets=[my_asset]) + def process(**context): + asset_store = context["asset_store"][my_asset] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +To see asset store in-action in a real DAG, checkout the DAG in `example_asset_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_asset_store.py>`_. + +Accessing asset store in a ``BaseEventTrigger`` +----------------------------------------------- + +When building Triggers used for asset "watching", asset store can be retrieved using the ``self.asset_store`` attribute. + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(name="my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +In the example above, ``my_data`` is created using the ``name`` However, the ``uri`` can also be used: + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(uri="s3://bucket/my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +Single-inlet shorthand +~~~~~~~~~~~~~~~~~~~~~~~ + +For tasks with exactly **one** concrete inlet, you can call ``get``, ``set``, ``delete``, and ``clear`` directly on ``context["asset_store"]`` without subscripting. + +.. code-block:: python + + @task(inlets=[my_asset], outlets=[my_asset]) + def process_single(**context): + asset_store = context["asset_store"] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +If the task has more than one concrete inlet, calling the shorthand raises a ``ValueError``. Use the subscript form (``context["asset_store"][my_asset]``) whenever a task has multiple inlets. Review Comment: We should also mention an example as to how it looks like when multiple inlets are present, aka the multi inlet pattern ########## airflow-core/docs/core-concepts/asset-store.rst: ########## @@ -0,0 +1,247 @@ + .. 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. + +.. _concepts:asset-store: + +.. spelling:word-list:: + + subscripted + subscripting + +Asset Store +=========== + +.. versionadded:: 3.3 + +Asset store is a persistent key/value store scoped to an *asset*, independent of any particular DAG run. Unlike :doc:`task store </core-concepts/task-store>`, which is tied to a single task instance, asset store persists across runs and is logically owned by the asset itself. It is the natural home for cross-run metadata such as watermarks, incremental-load cursors, and per-asset configuration. + +Asset store is accessed through the task context via ``context["asset_store"]``. + + +When is ``asset_store`` available? +------------------------------------ + +When using asset store within a task, ``context["asset_store"]`` is populated for **concrete** :class:`~airflow.sdk.definitions.asset.Asset` inlets and outlets. A task must declare at least one concrete inlet or outlet for ``asset_store`` to contain any entries. + +If using asset store in a ``BaseEventTrigger``, the ``self.asset_store`` parameter can be used within the ``BaseEventTrigger``. It can be subscripted in the same way that ``context["asset_store"]`` can be. + +.. warning:: + + **Outlets-only tasks**: if a task declares only ``outlets`` (no ``inlets``), ``context["asset_store"][my_asset]`` may raise a ``KeyError`` at runtime. The workaround is to declare the asset in **both** ``inlets`` and ``outlets``. + + .. code-block:: python + + # my_asset defined above ... + + + @task(inlets=[my_asset], outlets=[my_asset]) + def write_asset(**context): + context["asset_store"][my_asset].set("watermark", "2024-01-01") + + This known issue will be resolved in a future release. + + +Accessing asset store using ``context`` +--------------------------------------- + +An asset can be brought into "scope" (for lack of a better phrase) by including it in ``inlets`` (or both ``inlets`` and ``outlets``). Then subscript ``context["asset_store"]`` with the asset object to retrieve the asset store. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + my_asset = Asset("my_data", uri="s3://bucket/my_data") + + with DAG("example_asset_store", schedule=None): + + @task(inlets=[my_asset], outlets=[my_asset]) + def process(**context): + asset_store = context["asset_store"][my_asset] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +To see asset store in-action in a real DAG, checkout the DAG in `example_asset_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_asset_store.py>`_. + +Accessing asset store in a ``BaseEventTrigger`` +----------------------------------------------- + +When building Triggers used for asset "watching", asset store can be retrieved using the ``self.asset_store`` attribute. + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(name="my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +In the example above, ``my_data`` is created using the ``name`` However, the ``uri`` can also be used: + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(uri="s3://bucket/my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +Single-inlet shorthand +~~~~~~~~~~~~~~~~~~~~~~~ + +For tasks with exactly **one** concrete inlet, you can call ``get``, ``set``, ``delete``, and ``clear`` directly on ``context["asset_store"]`` without subscripting. + +.. code-block:: python + + @task(inlets=[my_asset], outlets=[my_asset]) + def process_single(**context): + asset_store = context["asset_store"] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +If the task has more than one concrete inlet, calling the shorthand raises a ``ValueError``. Use the subscript form (``context["asset_store"][my_asset]``) whenever a task has multiple inlets. + + +API reference +------------- + +The following methods are available on both the per-asset accessor (``context["asset_store"][my_asset]``), the shorthand (``context["asset_store"]``) when the task has exactly one inlet, and when using the ``self.asset_store`` attribute. + +``get(key)`` +~~~~~~~~~~~~ + +Returns the stored JSON value, or ``None`` if the key does not exist. + +.. code-block:: python + + # Using context + watermark = context["asset_store"][my_asset].get("watermark") + + # Using self.asset_store + watermark = self.asset_store[my_asset].get("watermark") + +``set(key, value)`` +~~~~~~~~~~~~~~~~~~~ + +Writes or overwrites a key-value pair. Unlike task store, asset store has no ``retention`` parameter. Values persist until explicitly deleted or until the asset is deactivated. Like with task store, ``value`` can be any JSON-compatible type, except for ``None``. This includes: + +* ``str`` +* ``int`` +* ``float`` +* ``bool`` +* ``list`` +* ``dict`` + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].set("watermark", "2024-06-01T00:00:00Z") + + # Using self.asset_store + self.asset_store[my_asset].set("watermark", "2024-06-01T00:00:00Z") + +``delete(key)`` +~~~~~~~~~~~~~~~ + +Deletes a single key. No-op if the key does not exist. + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].delete("watermark") + + # Using self.asset_store + self.asset_store[my_asset].delete("watermark") + +``clear()`` +~~~~~~~~~~~ + +Deletes *all* asset store keys for the asset. + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].clear() + + # Using self.asset_store + self.asset_store[my_asset].clear() + +Use cases +--------- + +Watermark pattern +~~~~~~~~~~~~~~~~~ + +The canonical use case for asset store is an incremental-load task that advances a watermark on each run. The watermark is stored on the asset itself so any task that reads or writes that asset can access it. This use case is especially applicable when building things like asset "watchers" using ``BaseEventTrigger``'s. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + orders = Asset("orders", uri="s3://data/orders/") + + + with DAG("incremental_orders", schedule="@daily"): + + @task(inlets=[orders], outlets=[orders]) + def load_new_orders(**context): + asset_store = context["asset_store"] # single-inlet shorthand + + # Read the last watermark, default to epoch if first run. + watermark = asset_store.get("watermark") or "1970-01-01T00:00:00Z" Review Comment: ```suggestion watermark = asset_store.get("watermark", default="1970-01-01T00:00:00Z") ``` ########## airflow-core/docs/administration-and-deployment/task-and-asset-store.rst: ########## @@ -0,0 +1,210 @@ + .. 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. + +.. _task-and-asset-store: + +Task and Asset Store Configuration +==================================== + +.. versionadded:: 3.3 + +The task and asset store is the persistence layer for :doc:`task store </core-concepts/task-store>` and :doc:`asset store </core-concepts/asset-store>`. By default, both are stored in the Airflow metadata database. This page describes the available configuration options, garbage-collection semantics, and how to provide a custom backend. + +Configuration reference +----------------------- + +All options live under the ``[state_store]`` section of ``airflow.cfg``. + +.. note:: + + The config section is ``[state_store]``, **not** ``[task_store]``. + +``backend`` +~~~~~~~~~~~ + +Full dotted path to a class that implements :class:`~airflow.sdk.state.BaseStateBackend`. Defaults to the built-in metastore backend. + +.. code-block:: ini + + [state_store] + backend = mypackage.state.CustomStateBackend + +``default_retention_days`` +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Number of days to retain **task store** rows after their last update. Rows older than this are deleted during the next garbage collection pass. + +* Set to ``0`` to disable time-based cleanup entirely. +* Default: ``30``. +* This setting does **not** apply to asset store rows. + +.. code-block:: ini + + [state_store] + default_retention_days = 30 + +``clear_on_success`` +~~~~~~~~~~~~~~~~~~~~ + +When ``True``, all task store keys for a task instance are automatically deleted when that task instance moves to the ``success`` state. Defaults to ``False``, which preserves task store entries after success for observability (e.g. the submitted job ID or the last row count is still readable from the UI or REST API after the run completes). + +.. important:: + + ``clear_on_success`` clears **task store only**. It has no effect on asset store. Asset store is scoped to the asset rather than the task instance and must be cleared explicitly. + +.. code-block:: ini + + [state_store] + clear_on_success = False + +``state_cleanup_batch_size`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Number of rows deleted per batch during garbage collection cleanup. Set to ``0`` (default) to delete all matching rows in a single statement. Tune this on deployments with large ``task_store`` tables to reduce lock contention. + +.. code-block:: ini + + [state_store] + state_cleanup_batch_size = 10000 + +.. _task-and-asset-store:worker-backends: + +Worker-side backend (``[workers] state_backend``) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +A separate, optional config key under ``[workers]`` lets you route task store and asset store values through a worker-side backend before they reach the API server. + +.. code-block:: ini + + [workers] + state_backend = mypackage.state.S3StateBackend + +When this is set, ``TaskStoreAccessor.set()`` calls ``serialize_task_store_to_ref()`` on the worker-side backend before sending the returned value (a reference to the actual storage) to the Execution API, and ``get()`` calls ``deserialize_task_store_from_ref()`` after receiving the stored reference from the Execution API. See `Custom worker-side backends`_ below. + + +Garbage collection semantics +----------------------------- + +The cleanup task, also known as "garbage collection" is triggered using the Airflow CLI. The command to trigger the cleanup task is ``airflow state-store cleanup-task-store``. This process removes store rows according to the following rules: + +**Time-based expiry (task store only)** + Rows whose ``expires_at < now()`` are deleted. ``expires_at`` is computed on the *worker* at write time, not by the server. + +**``default_retention_days`` fallback (task store only)** + Keys written with no explicit ``retention`` (i.e. ``expires_at`` is ``NULL``) are governed by the global ``default_retention_days`` setting. When this setting is positive, the garbage collection job treats such rows as expiring ``default_retention_days`` days after their last update. + +**``NEVER_EXPIRE`` keys** + Keys set with ``retention=NEVER_EXPIRE`` are stored with ``expires_at = NULL`` and a flag that tells the garbage collection to skip them unconditionally. They are never deleted by time-based cleanup, regardless of ``default_retention_days``. + +**Orphan sweep (asset store)** + Asset store rows for assets that no longer have an ``asset_active`` record are deleted during the orphan-sweep pass. This cleans up store entries for deactivated or renamed assets. + +.. important:: + + Garbage collection only works for the ``MetastoreStateBackend``. Custom backends are explicitly skipped. + + + +Custom backends +--------------- + +A custom backend must subclass :class:`~airflow.sdk.state.BaseStateBackend` and implement its abstract methods: ``get``, ``set``, ``delete``, and ``clear`` for synchronous callers and the ``aget``, ``aset``, ``adelete``, and ``aclear`` async equivalents. Refer to :class:`~airflow.sdk.state.BaseStateBackend` for the full API. + +Each method receives a ``scope`` argument that is either a :class:`~airflow.sdk.state.TaskScope` or an :class:`~airflow.sdk.state.AssetScope`. Use ``isinstance`` to dispatch: + +.. code-block:: python + + from airflow.sdk.state import BaseStateBackend, TaskScope, AssetScope + + + class MyBackend(BaseStateBackend): + def get(self, scope, key, *, session=None): + if isinstance(scope, TaskScope): + return self._task_store.get(scope, key) + elif isinstance(scope, AssetScope): + return self._asset_store.get(scope, key) + +:class:`~airflow.sdk.state.AssetScope` has three optional fields: ``asset_id`` (integer, server-side only), ``name``, and ``uri``. At least one must be set. Server-side operations (REST API calls) provide ``asset_id``. Worker-side operations provide ``name`` or ``uri`` (workers do not have access to the integer ``asset_id``). + +Configure the class via ``[state_store] backend``: + +.. code-block:: ini + + [state_store] + backend = mypackage.state.MyBackend + + +Custom worker-side backends +---------------------------- + +Worker-side backends extend ``BaseStateBackend`` with two pairs of serialization hooks. They are configured separately via ``[workers] state_backend`` and run *on the worker process*, not on the API server. This lets you store large payloads or credentialed data directly using worker infrastructure while only a compact reference string is kept in the database. + +Override four serialization hooks from :class:`~airflow.sdk.state.BaseStateBackend`: + +* ``serialize_task_store_to_ref``: called by ``TaskStoreAccessor.set()`` before the value is sent to the Execution API; return a compact reference string (e.g. an S3 key) to be stored in the database instead of the raw value. +* ``deserialize_task_store_from_ref``: called by ``TaskStoreAccessor.get()`` after retrieving the reference; return the actual value. Review Comment: ```suggestion * ``deserialize_task_store_from_ref``: called by ``TaskStoreAccessor.get()`` after retrieving the reference from the backend; return the actual value. ``` ########## airflow-core/docs/core-concepts/asset-store.rst: ########## @@ -0,0 +1,247 @@ + .. 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. + +.. _concepts:asset-store: + +.. spelling:word-list:: + + subscripted + subscripting + +Asset Store +=========== + +.. versionadded:: 3.3 + +Asset store is a persistent key/value store scoped to an *asset*, independent of any particular DAG run. Unlike :doc:`task store </core-concepts/task-store>`, which is tied to a single task instance, asset store persists across runs and is logically owned by the asset itself. It is the natural home for cross-run metadata such as watermarks, incremental-load cursors, and per-asset configuration. + +Asset store is accessed through the task context via ``context["asset_store"]``. + + +When is ``asset_store`` available? +------------------------------------ + +When using asset store within a task, ``context["asset_store"]`` is populated for **concrete** :class:`~airflow.sdk.definitions.asset.Asset` inlets and outlets. A task must declare at least one concrete inlet or outlet for ``asset_store`` to contain any entries. + +If using asset store in a ``BaseEventTrigger``, the ``self.asset_store`` parameter can be used within the ``BaseEventTrigger``. It can be subscripted in the same way that ``context["asset_store"]`` can be. + +.. warning:: + + **Outlets-only tasks**: if a task declares only ``outlets`` (no ``inlets``), ``context["asset_store"][my_asset]`` may raise a ``KeyError`` at runtime. The workaround is to declare the asset in **both** ``inlets`` and ``outlets``. + + .. code-block:: python + + # my_asset defined above ... + + + @task(inlets=[my_asset], outlets=[my_asset]) + def write_asset(**context): + context["asset_store"][my_asset].set("watermark", "2024-01-01") + + This known issue will be resolved in a future release. + + +Accessing asset store using ``context`` +--------------------------------------- + +An asset can be brought into "scope" (for lack of a better phrase) by including it in ``inlets`` (or both ``inlets`` and ``outlets``). Then subscript ``context["asset_store"]`` with the asset object to retrieve the asset store. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + my_asset = Asset("my_data", uri="s3://bucket/my_data") + + with DAG("example_asset_store", schedule=None): + + @task(inlets=[my_asset], outlets=[my_asset]) + def process(**context): + asset_store = context["asset_store"][my_asset] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +To see asset store in-action in a real DAG, checkout the DAG in `example_asset_store.py <https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/example_dags/example_asset_store.py>`_. + +Accessing asset store in a ``BaseEventTrigger`` +----------------------------------------------- + +When building Triggers used for asset "watching", asset store can be retrieved using the ``self.asset_store`` attribute. + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(name="my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +In the example above, ``my_data`` is created using the ``name`` However, the ``uri`` can also be used: + +.. code-block:: python + + from airflow.sdk import Asset, BaseEventTrigger, TriggerEvent + from collections.abc import AsyncIterator + + + class GenericEventTrigger(BaseEventTrigger): + ... + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Logic that fires a TriggerEvent.""" + my_data = Asset(uri="s3://bucket/my_data") + asset_store = self.asset_store[my_data] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +Single-inlet shorthand +~~~~~~~~~~~~~~~~~~~~~~~ + +For tasks with exactly **one** concrete inlet, you can call ``get``, ``set``, ``delete``, and ``clear`` directly on ``context["asset_store"]`` without subscripting. + +.. code-block:: python + + @task(inlets=[my_asset], outlets=[my_asset]) + def process_single(**context): + asset_store = context["asset_store"] + watermark = asset_store.get("watermark") + asset_store.set("watermark", "2024-06-01") + +If the task has more than one concrete inlet, calling the shorthand raises a ``ValueError``. Use the subscript form (``context["asset_store"][my_asset]``) whenever a task has multiple inlets. + + +API reference +------------- + +The following methods are available on both the per-asset accessor (``context["asset_store"][my_asset]``), the shorthand (``context["asset_store"]``) when the task has exactly one inlet, and when using the ``self.asset_store`` attribute. + +``get(key)`` +~~~~~~~~~~~~ + +Returns the stored JSON value, or ``None`` if the key does not exist. + +.. code-block:: python + + # Using context + watermark = context["asset_store"][my_asset].get("watermark") + + # Using self.asset_store + watermark = self.asset_store[my_asset].get("watermark") + +``set(key, value)`` +~~~~~~~~~~~~~~~~~~~ + +Writes or overwrites a key-value pair. Unlike task store, asset store has no ``retention`` parameter. Values persist until explicitly deleted or until the asset is deactivated. Like with task store, ``value`` can be any JSON-compatible type, except for ``None``. This includes: + +* ``str`` +* ``int`` +* ``float`` +* ``bool`` +* ``list`` +* ``dict`` + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].set("watermark", "2024-06-01T00:00:00Z") + + # Using self.asset_store + self.asset_store[my_asset].set("watermark", "2024-06-01T00:00:00Z") + +``delete(key)`` +~~~~~~~~~~~~~~~ + +Deletes a single key. No-op if the key does not exist. + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].delete("watermark") + + # Using self.asset_store + self.asset_store[my_asset].delete("watermark") + +``clear()`` +~~~~~~~~~~~ + +Deletes *all* asset store keys for the asset. + +.. code-block:: python + + # Using context + context["asset_store"][my_asset].clear() + + # Using self.asset_store + self.asset_store[my_asset].clear() + +Use cases +--------- + +Watermark pattern +~~~~~~~~~~~~~~~~~ + +The canonical use case for asset store is an incremental-load task that advances a watermark on each run. The watermark is stored on the asset itself so any task that reads or writes that asset can access it. This use case is especially applicable when building things like asset "watchers" using ``BaseEventTrigger``'s. + +.. code-block:: python + + from airflow.sdk import Asset, DAG, task + + orders = Asset("orders", uri="s3://data/orders/") + + Review Comment: ```suggestion ``` -- 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]
