[WIP] AIP-93 Asset Watermarks and State Variables

Status

Motivation

Incremental processing is a very common pattern in Airflow use cases. In some pipelines, this is strictly necessary. Although external event driven scheduling is supported today in Airflow through AIP-82, incremental polling is not something that has been easy to implement. In more traditional DAG Authoring (Sensors, Operators, etc.), users are forced to use XCom or Variables for storing values (watermarks) for incremental processing, or inventing and implementing their own solution. In the case of event-driven triggering, it’s even more challenging for users to build their custom incremental processing implantation, due to the lack of XCom and Variable support. 

There have been several attempts to store state within a child of the BaseEventTrigger to manage incremental processes, while none of these have proven to be effective or robust. This has seriously hampered the ability for the community to build logic to monitor Assets such as object stores, SQL databases, and other non-queue/stream-based Assets.

For Triggers built for Asset “watching”, it is helpful, if not essential to persist some state.

Considerations

What change do you propose to make?

At a high level, we propose to make the following changes:

• Introducing a StateVariable model for persisting Asset/Trigger state metadata. This would be a general model, not limited to event-driven scheduling, allowing for Airflow users to store and retrieve state as specified in EventTrigger code, as well as within things like Operators, Sensors, etc.
• This model should be reusable and extendable in other Airflow components (Workers, Scheduler, etc). The way it fetches and stores the data should also vary depending on the components.
• (Optional) Adding some sort of helper and/or updating BaseEventTrigger to include methods allowing users to store and retrieve values (watermarks). 
• Create a pluggable backend, similar to the XCom model, to act as a state store for users not interested in using the “batteries included” model.
• Adding a UI component to surface StateVariable 's.

What problem does it solve?

This solves one of the most glaring problems with building Triggers compatible with "Asset-watching", or event-driven scheduling; the inability to persist state. Without a model like StateVariable , it's nearly impossible to do things like monitor an S3 bucket for new files landing or handle the addition/removal of a new row to a SQL database. Despite being one of the most touted features of Airflow 3.0, the community has been very slow to develop and distribute Event Triggers to be used for "Asset-watching". This AIP aims to address this problem in an elegant and "Airflow-onic" way, using a "batteries included" model, along with a pluggable backend.

Technical Details

Stage 1 - Building the generic StateVariable model

Building the async-aware StateVariable model would require minimal lift. Like the Variable model, this would be “unrelated” to other models, making the creation a bit easier. Such model should be usable by various Airflow components, other than triggers. Therefore, it should provide an abstraction layer that 

This model contains two core fields: key and value. Below is the pseudo-code for the StateVariable model. 

class StateVariable:
    """
    StateVariable class allows various methods fetching and storing values.
    """
	key: str
	value: str

    def set(self, key, value):
        # Using db connection as an example here, this method will vary
        db_backend.set_value(key=key, value=value)

    def get(self):
        return db_backend.get_value(key=key)

Despite their similarities, StateVariable will differ from Variable in a number of ways. These include:

• StateVariable will NOT use a traditional Secrets Backend. Instead, it will have its own type of backend.
• StateVariable will have its own UI component(s), and can be integrated more "natively" into a DAG/Asset Watcher.
• StateVariable will have a sort of "owner" field, something that Variable  does not have.

Stage 2 - Extending the backend of StateVariable

The state model will be reusable and extendable in other Airflow components, such as task, worker, scheduler, etc. It should support not only the API, but also the Task SDK for fetching and storing the value for different components.

Stage 3 - Implement the state-enabled Trigger/Asset Watcher

Once the StateVariable  model has been created, it can be used within a BaseEventTrigger  just like Variable  would; using StateVariable.get()  and StateVariable.set(...) . This pattern is quite intuitive for DAG authors.

It would require some sort of get and set logic to be implemented in a class that inherits from BaseEventTrigger in order to store and retrieve watermark value. This could be used in a standalone manner, or be implemented in the form of a "helper" method in the BaseEventTrigger class.

Examples

“Watching” an S3 Bucket

One of the most common use-cases for event-driven scheduling will most likely be “watching” an object store for changes. Ideally, each time that an AssetWatcher’s Trigger runs, it will not re-scan the entire bucket. The workflow would look something like this:

1. After the first Trigger run, a timestamp of this “scan” for data is recorded.
2. On the second run, the Trigger retrieves the timestamp (watermark) from the completion of the last run and only scans the object store for files after that watermark.
3. Trigger writes the watermark again.
4. Process repeats for N Trigger runs following.

Incremental Data in SQL Databases

Another common use-case for event-driven scheduling is “watching” tables in relational databases; we’ll use Postgres as an example. Outside of traditional CDC, it’s common for data teams to use an updated column in a Postgres database to upsert data. Using the StateVariable approach, Airflow users can use the following workflow. This should feel quite similar to the workflow for “watching” an S3 Bucket; that’s intentional, as we’re trying to implement a repeatable pattern.

1. After the first Trigger run, a timestamp of this “scan” for data is recorded and stored to the table_last_updated_at StateVariable.
2. On the second run, the Trigger retrieves the watermark (value) from the completion of the last run and only scans the updated column in the desired table for records with a timestamp greater than the watermark.
3. Trigger writes the watermark again.
4. Process repeats for N Trigger runs following.

Below is an example of a Trigger that executes a SQL query to incrementally read data, with a watermark. Implementing something like the new_watermark  and get_watermark  methods is not required; it just shows an example of how StateVariable  can be used.

from airflow.sdk import StateVariable

class SQLIncrementalQueryTrigger(BaseEventTrigger):
	state_variable_name = "table_last_updated_at"
    ...
	
	@property
	def new_watermark(self):
		return datetime.now()
	
	@property
	def get_watermark(self):
		return StateVariable.get(self.state_variable_name)

    async def run(self):
        while True:
            results = conn.get_records("SELECT * FROM table WHERE updated_at > '%s'", self.get_watermark())

            if results:
    		 	StateVariable.set(key=self.state_variable_name, value=self.new_watermark)
                yield TriggerEvent({"status": "success", "results": results})
                break
            else:
                await asyncio.sleep(self.poke_interval)

Why is it needed?

Allow for an Airflow-supported methodology for persisting state and unblocking the development of Event Triggers for Airflow users interested in "Asset-watching".

Are there any downsides to this change?

There is one theorized downside; Airflow users leveraging the StateVariable  model as a sort of key-value store outside of an Event Trigger. However, this is not something that is overly concerning. It would be much more likely for Airflow users to abuse the Variable  model instead. Along with this are the general downsides to added complexity, but those are minimal with this proposal.

Which users are affected by the change?

DAG Authors: these users will now have access to a tool that makes authoring Triggers, Sensors, and Tasks used to orchestrate incremental processes more accessible and unlocks the ability to further build out event-driven logic.

Deployment Managers: introducing a pattern with a pluggable backend may provide deployment managers an additional component of their Airflow “stack” to stand up and manage. This will also require a DB migration, which is outlined below.

How are users affected by the change? (e.g. DB upgrade required?)

Current users who don't use incremental event triggers are unaffected by this change. However, this new feature requires a DB migration, since a new model is being added. The only changes needed in this case would be the addition of a new table.

What is the level of migration effort (manual and automated) needed for the users to adapt to the breaking changes? (especially in context of Airflow 3)

There are not breaking changes included as part of this AIP. Added the StateVariable model will require a DB migration, but forces no breaking changes upon users.

Other considerations?

• This would be compatible with Asset partitions.
• For now, this sort of “state store” would be limited to Asset “watching”.
• Naming is going to be important, both for the state store, as well as the properties/methods used to store/retrieve state. Some options that have been throw around include State , StateVariable , ProcessState , Watermark , and more.
• One of the biggest challenges that we're facing is how this is tied to an Asset. A Trigger defined for an AssetWatcher  is in no way tied to the Asset itself. This disconnect makes the naming and implementation of this functionality challenging. Should this be something that is tied to an Asset? Should it be only applicable at the Trigger-level? Overall, the goal of this AIP is to make event-driven scheduling more intuitive for Airflow users. 
• One of the potential future works is to add StateVariable to the Airflow UI, so users can see exactly the StateVariable values, similar to how Airflow Variable is presented.

What defines this AIP as "done"?

This AIP will be considered “done” when the PR creates the model and the pluggable backend is merged and the solution is well-documented.