Add asset state store commands to airflowctl

Author: justinpakzad

airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py

@@ -56,6 +56,11 @@ def date_param():
     "assets list",
     "assets get 1",
     "assets create-event --asset-id=1",
+    "assets set-state-store 1 test_key test_value",
+    "assets get-state-store 1 test_key",
+    "assets list-state-store 1",
+    "assets delete-state-store 1 test_key",
+    "assets clear-state-store 1",
     # Backfill commands
     "backfill list example_bash_operator",
     # Config commands

airflow-ctl/docs/images/command_hashes.txt

@@ -1,5 +1,5 @@
 main:2c358f5fc894541cab11854dcd67658a
-assets:f225579fbecb3d919d152b4f1ba83185
+assets:949f2c4ff70e32f46c905931e82dcb56
 auth:d79e9c7d00c432bdbcbc2a86e2e32053
 backfill:74c8737b0a62a86ed3605fa9e6165874
 config:a3d936cb15fe3b547bf6c82cf93d923f

airflow-ctl/docs/images/output_assets.svg

@@ -32,6 +32,9 @@
     AssetCollectionResponse,
     AssetEventResponse,
     AssetResponse,
+    AssetStateStoreBody,
+    AssetStateStoreCollectionResponse,
+    AssetStateStoreResponse,
     BackfillCollectionResponse,
     BackfillPostBody,
     BackfillResponse,
@@ -326,6 +329,39 @@ def delete_queued_event(self, dag_id: str, asset_id: str) -> str | ServerRespons
         self.client.delete(f"dags/{dag_id}/assets/{asset_id}/queuedEvents")
         return asset_id
 
+    def list_state_store(self, asset_id: str) -> AssetStateStoreCollectionResponse | ServerResponseError:
+        """List all state store entries for an asset."""
+        return super().execute_list(
+            path=f"/assets/{asset_id}/state-store", data_model=AssetStateStoreCollectionResponse
+        )
+
+    def get_state_store(self, asset_id: str, key: str) -> AssetStateStoreResponse | ServerResponseError:
+        """Get a single asset state store entry."""
+        self.response = self.client.get(f"assets/{asset_id}/state-store/{key}")
+        return AssetStateStoreResponse.model_validate_json(self.response.content)
+
+    def set_state_store(self, asset_id: str, key: str, value: str) -> str | ServerResponseError:
+        """Set an asset state store value. Creates or overwrites the key."""
+        try:
+            parsed_value = json.loads(value)
+        except (ValueError, TypeError):
+            parsed_value = value
+        self.client.put(
+            f"assets/{asset_id}/state-store/{key}",
+            json=AssetStateStoreBody(value=parsed_value).model_dump(mode="json"),
+        )
+        return key
+
+    def delete_state_store(self, asset_id: str, key: str) -> str | ServerResponseError:
+        """Delete a single asset state store key."""
+        self.client.delete(f"assets/{asset_id}/state-store/{key}")
+        return key
+
+    def clear_state_store(self, asset_id: str) -> str | ServerResponseError:
+        """Delete all state store keys for an asset."""
+        self.client.delete(f"assets/{asset_id}/state-store")
+        return asset_id
+
 
 class BackfillOperations(BaseOperations):
     """Backfill operations."""

airflow-ctl/src/airflowctl/api/operations.py

@@ -510,6 +510,7 @@ def __init__(self, file_path: str | Path | None = None):
             "trigger",
             "add",
             "edit",
+            "set",
             "clear",
         ]
         # Datamodels whose generated bool flags follow the datamodel field defaults instead of

airflow-ctl/src/airflowctl/ctl/cli_config.py

@@ -28,6 +28,11 @@ assets:
   delete-queued-events: "Delete all queued events for a given asset"
   delete-dag-queued-events: "Delete all queued asset events for a given Dag"
   delete-queued-event: "Delete a specific queued asset event for a given Dag and asset"
+  list-state-store: "List all state store entries for an asset"
+  get-state-store: "Retrieve a single asset state store entry by its ID and key"
+  set-state-store: "Set an asset state store value. Creates or overwrites the key."
+  delete-state-store: "Delete a single asset state store key by its ID and key"
+  clear-state-store: "Delete all state store keys for an asset by its ID"
 
 backfill:
   create: "Create a backfill job for a given Dag ID and date range"

airflow-ctl/src/airflowctl/ctl/help_texts.yaml

@@ -36,6 +36,8 @@
     AssetCollectionResponse,
     AssetEventResponse,
     AssetResponse,
+    AssetStateStoreCollectionResponse,
+    AssetStateStoreResponse,
     BackfillCollectionResponse,
     BackfillPostBody,
     BackfillResponse,
@@ -321,7 +323,11 @@ class TestAssetsOperations:
         queued_events=[asset_queued_event_response],
         total_entries=1,
     )
-
+    asset_state_store_response = AssetStateStoreResponse(
+        key="my_key",
+        value={"my_val": 0},  # type: ignore[arg-type]
+        updated_at=datetime.datetime(2025, 1, 1, 0, 0, 0),
+    )
     dag_run_response = DAGRunResponse(
         dag_display_name=dag_id,
         dag_run_id=dag_id,
@@ -513,6 +519,73 @@ def handle_request(request: httpx.Request) -> httpx.Response:
         response = client.assets.delete_queued_event(dag_id=self.dag_id, asset_id=self.asset_id)
         assert response == self.asset_id
 
+    def test_list_state_store(self):
+        collection_response = AssetStateStoreCollectionResponse(
+            asset_state_store=[self.asset_state_store_response],
+            total_entries=1,
+        )
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.url.path == f"/api/v2/assets/{self.asset_id}/state-store"
+            return httpx.Response(200, json=json.loads(collection_response.model_dump_json()))
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.assets.list_state_store(self.asset_id)
+        assert response == collection_response
+
+    def test_get_state_store(self):
+        key = self.asset_state_store_response.key
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.url.path == f"/api/v2/assets/{self.asset_id}/state-store/{key}"
+            return httpx.Response(200, json=json.loads(self.asset_state_store_response.model_dump_json()))
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.assets.get_state_store(self.asset_id, key)
+        assert response == self.asset_state_store_response
+
+    @pytest.mark.parametrize(
+        ("value", "expected"),
+        [
+            ('{"index": 0}', {"index": 0}),
+            ("hello", "hello"),
+        ],
+    )
+    def test_set_state_store(self, value, expected):
+        key = self.asset_state_store_response.key
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.method == "PUT"
+            assert request.url.path == f"/api/v2/assets/{self.asset_id}/state-store/{key}"
+            assert json.loads(request.content) == {"value": expected}
+            return httpx.Response(204)
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.assets.set_state_store(self.asset_id, key, value)
+        assert response == key
+
+    def test_delete_state_store(self):
+        key = self.asset_state_store_response.key
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.method == "DELETE"
+            assert request.url.path == f"/api/v2/assets/{self.asset_id}/state-store/{key}"
+            return httpx.Response(204)
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.assets.delete_state_store(self.asset_id, key)
+        assert response == key
+
+    def test_clear_state_store(self):
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            assert request.method == "DELETE"
+            assert request.url.path == f"/api/v2/assets/{self.asset_id}/state-store"
+            return httpx.Response(204)
+
+        client = make_api_client(transport=httpx.MockTransport(handle_request))
+        response = client.assets.clear_state_store(self.asset_id)
+        assert response == self.asset_id
+
 
 class TestBackfillOperations:
     backfill_id: NonNegativeInt = 1