From 9b8ceb4af8e3fbc5308648848c191a7ee88c963e Mon Sep 17 00:00:00 2001
From: jbirddog <100367399+jbirddog@users.noreply.github.com>
Date: Mon, 10 Jun 2024 11:17:41 -0400
Subject: [PATCH] Prevent KeyError for optional data store writes (#1700)

---
 .../src/spiffworkflow_backend/data_stores/json.py             | 4 ++++
 .../src/spiffworkflow_backend/data_stores/kkv.py              | 2 ++
 .../src/spiffworkflow_backend/data_stores/typeahead.py        | 2 ++
 3 files changed, 8 insertions(+)

diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/json.py b/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/json.py
index ff141f473..4fef62d2a 100644
--- a/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/json.py
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/json.py
@@ -69,6 +69,8 @@ class JSONDataStore(BpmnDataStoreSpecification, DataStoreCRUD):  # type: ignore
 
     def set(self, my_task: SpiffTask) -> None:
         """set."""
+        if self.bpmn_id not in my_task.data:
+            return
         model: JSONDataStoreModel | None = None
         location = self.data_store_location_for_task(JSONDataStoreModel, my_task, self.bpmn_id)
 
@@ -127,6 +129,8 @@ class JSONFileDataStore(BpmnDataStoreSpecification):  # type: ignore
 
     def set(self, my_task: SpiffTask) -> None:
         """set."""
+        if self.bpmn_id not in my_task.data:
+            return
         location = self._data_store_location_for_task(my_task, self.bpmn_id)
         if location is None:
             raise DataStoreWriteError(f"Unable to write to data store '{self.bpmn_id}' using location '{location}'.")
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/kkv.py b/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/kkv.py
index f7e85cac7..0b47de635 100644
--- a/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/kkv.py
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/kkv.py
@@ -114,6 +114,8 @@ class KKVDataStore(BpmnDataStoreSpecification, DataStoreCRUD):  # type: ignore
         my_task.data[self.bpmn_id] = getter
 
     def set(self, my_task: SpiffTask) -> None:
+        if self.bpmn_id not in my_task.data:
+            return
         location = self.data_store_location_for_task(KKVDataStoreModel, my_task, self.bpmn_id)
         store_model: KKVDataStoreModel | None = None
 
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/typeahead.py b/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/typeahead.py
index 3d5ad4d93..f92626f1e 100644
--- a/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/typeahead.py
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/data_stores/typeahead.py
@@ -43,6 +43,8 @@ class TypeaheadDataStore(BpmnDataStoreSpecification, DataStoreCRUD):  # type: ig
 
     def set(self, my_task: SpiffTask) -> None:
         """set."""
+        if self.bpmn_id not in my_task.data:
+            return
         typeahead_data_by_category = my_task.data[self.bpmn_id]
         for category, items in typeahead_data_by_category.items():
             db.session.query(TypeaheadModel).filter_by(category=category).delete()