AIP-103: Allow for get/set/delete/clear to run for AssetStateStoreAccessor on Triggerer

[jroachgolf84]

Description

When attempting to use the Asset State Store from within a BaseEventTrigger, the following exception was thrown:

Traceback (most recent call last):
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/supervisor.py", line 890, in handle_requests
msg = self.decoder.validate_python(self._deserialize_request(request.body))
File "/usr/python/lib/python3.10/site-packages/pydantic/type_adapter.py", line 441, in validate_python
return self.validator.validate_python(
pydantic_core._pydantic_core.ValidationError: 1 validation error for tagged-union[
TriggerStateChanges, GetConnection, DeleteVariable, GetVariable, GetVariableKeys,
PutVariable, DeleteXCom, GetXCom, SetXCom, GetTICount, GetTaskStates, GetDagRunState,
GetDRCount, GetPreviousTI, GetHITLDetailResponse, UpdateHITLDetail, MaskSecret
]
Input tag 'GetAssetStateStoreByName' found using 'type' does not match any of the expected tags:
'TriggerStateChanges', 'GetConnection', 'DeleteVariable', 'GetVariable', 'GetVariableKeys',
'PutVariable', 'DeleteXCom', 'GetXCom', 'SetXCom', 'GetTICount', 'GetTaskStates',
'GetDagRunState', 'GetDRCount', 'GetPreviousTI', 'GetHITLDetailResponse', 'UpdateHITLDetail',
'MaskSecret'
[type=union_tag_invalid, input_value={'name': 'generic_asset', ... 'type': 'GetAssetStateStoreByName'}, input_type=dict]

This was called out as something that needed to be addressed by @vikramkoka in #67839. This PR addresses this error by adding a pathway for each of the 8 operations that the Asset State Store can perform.

related: #67839

Testing

Changes were tested using both unit-tests, as well as E2E. Unit-tests can be run using the below commands:

breeze testing core-tests airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py
breeze testing core-tests airflow-core/tests/unit/jobs/test_triggerer_job.py

E2E Testing Example

To test this E2E, a custom BaseEventTrigger was written to test get, set, delete, and clear operations. The DAG that uses this Trigger executes successfully E2E, and the Triggered contains all of the expected logging statements.

fromairflow.triggers.baseimportBaseEventTrigger, TriggerEventfromcollections.abcimportAsyncIteratorfromtypingimportAnyimportasyncioimportloggingimportrandomclassGenericEventTrigger(BaseEventTrigger):
def__init__(
self,
random_number,
waiter_delay,
asset_name,
**kwargs
):
super().__init__(**kwargs)
self.random_number=random_numberself.waiter_delay=waiter_delayself.asset_name=asset_namedefserialize(self) ->tuple[str, dict[str, Any]]:
"""Serialize the Trigger, including the func, params, and waiter_delay."""return (
self.__class__.__module__+"."+self.__class__.__qualname__,
{
"random_number": self.random_number,
"waiter_delay": self.waiter_delay,
"asset_name": self.asset_name,
},
)
asyncdefrun(self) ->AsyncIterator[TriggerEvent]:
"""Logic that fires a TriggerEvent."""asset_state_store=self.asset_state_storelogging.info(f"asset_state_store: {asset_state_store}")
# Clearing state store before starting (in case there is anything there)asset_state_store.clear()
logging.info(f"cleared asset_state_store")
whileTrue:
# Get the previous numberlast_result=asset_state_store.get("result", -1)
logging.info(f"last_result: {last_result}")
# Create a new numberresult=random.randint(0, 5)
asset_state_store.set("result", result)
logging.info(f"result: {result}")
ifresult==self.random_number:
logging.info("yield'ing TriggerEvent")
yieldTriggerEvent({"status": "success", "result": result})
# Clear the state_storeasset_state_store.delete(key="result") # Should this work?logging.info(f"deleted 'result' key for asset_state_store")
breaklogging.info(f"Sleeping for {self.waiter_delay} seconds")
awaitasyncio.sleep(self.waiter_delay)

[amoghrajesh]

Some more work is needed.

[jroachgolf84]

@amoghrajesh - changes have been implemented.

[amoghrajesh]

This PR makes self.asset_state_store actually functional inside BaseEventTrigger.run(), so it is worth documenting on the asset-state-store.rst page.

Right now theres nothing explaining how to use self.asset_state_store from a trigger, that the triggerer injects it before run() is called. The example in the PR description is a good starting point.

[jroachgolf84]

@amoghrajesh - all feedback has been implemented!

[amoghrajesh]

Thanks for working on this @jroachgolf84! Looks good.

[amoghrajesh]

Actually let me just apply the suggestions.

[github-actions]

Backport successfully created: v3-3-test

Note: As of Merging PRs targeted for Airflow 3.X
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.

In matter of doubt please ask in #release-management Slack channel.

Status	Branch	Result
✅	v3-3-test