From a33475da648b08350ea8ff122225e0fdf58e41af Mon Sep 17 00:00:00 2001
From: lwolczynski <54366429+lwolczynski@users.noreply.github.com>
Date: Fri, 6 Dec 2024 11:05:49 -0600
Subject: [PATCH] IWF-372: Fix GetSearchAttributes returning nil when no
 IwfExecutionStateIds are set (#505)

---
 integ/wait_until_search_attributes_optimization_test.go  | 9 +++++----
 .../wait_until_search_attributes_optimization/routers.go | 6 ++++++
 service/common/mapper/searchAttribute.go                 | 8 ++++++++
 service/interpreter/stateExecutionCounter.go             | 1 +
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/integ/wait_until_search_attributes_optimization_test.go b/integ/wait_until_search_attributes_optimization_test.go
index d385c8ee..7e0f9038 100644
--- a/integ/wait_until_search_attributes_optimization_test.go
+++ b/integ/wait_until_search_attributes_optimization_test.go
@@ -140,18 +140,19 @@ func doTestWaitUntilHistoryCompleted(
 		assertions.Equal([]string{"S3", "S6"}, historyEventSAs(upsertSAEvents[7]))
 		assertions.Equal([]string{"S3"}, historyEventSAs(upsertSAEvents[8]))
 		assertions.Equal([]string{"null"}, historyEventSAs(upsertSAEvents[9]))
+	case iwfidl.DISABLED:
+		assertions.Equal(1, len(upsertSAEvents))
 	case iwfidl.ENABLED_FOR_STATES_WITH_WAIT_UNTIL:
+	default:
 		assertions.Equal(9, len(upsertSAEvents))
 		assertions.Equal([]string{"S1"}, historyEventSAs(upsertSAEvents[1]))
-		assertions.Equal([]string{"S2"}, historyEventSAs(upsertSAEvents[2]))
-		assertions.Equal([]string{"S2", "S3"}, historyEventSAs(upsertSAEvents[3]))
+		assertions.Equal([]string{"null"}, historyEventSAs(upsertSAEvents[2]))
+		assertions.Equal([]string{"S3"}, historyEventSAs(upsertSAEvents[3]))
 		assertions.Equal([]string{"S3", "S4"}, historyEventSAs(upsertSAEvents[4]))
 		assertions.Equal([]string{"S3"}, historyEventSAs(upsertSAEvents[5]))
 		assertions.Equal([]string{"S3", "S6"}, historyEventSAs(upsertSAEvents[6]))
 		assertions.Equal([]string{"S3"}, historyEventSAs(upsertSAEvents[7]))
 		assertions.Equal([]string{"null"}, historyEventSAs(upsertSAEvents[8]))
-	case iwfidl.DISABLED:
-		assertions.Equal(1, len(upsertSAEvents))
 	}
 }
 
diff --git a/integ/workflow/wait_until_search_attributes_optimization/routers.go b/integ/workflow/wait_until_search_attributes_optimization/routers.go
index 5ab4841b..b7a2230b 100644
--- a/integ/workflow/wait_until_search_attributes_optimization/routers.go
+++ b/integ/workflow/wait_until_search_attributes_optimization/routers.go
@@ -154,6 +154,9 @@ func (h *handler) ApiV1WorkflowStateDecide(c *gin.Context) {
 						NextStates: []iwfidl.StateMovement{
 							{
 								StateId: State2,
+								StateOptions: &iwfidl.WorkflowStateOptions{
+									SkipWaitUntil: iwfidl.PtrBool(true),
+								},
 							},
 						},
 					},
@@ -192,6 +195,9 @@ func (h *handler) ApiV1WorkflowStateDecide(c *gin.Context) {
 						NextStates: []iwfidl.StateMovement{
 							{
 								StateId: State2,
+								StateOptions: &iwfidl.WorkflowStateOptions{
+									SkipWaitUntil: iwfidl.PtrBool(true),
+								},
 							},
 							{
 								StateId: State3,
diff --git a/service/common/mapper/searchAttribute.go b/service/common/mapper/searchAttribute.go
index fba189d4..539ccf5a 100644
--- a/service/common/mapper/searchAttribute.go
+++ b/service/common/mapper/searchAttribute.go
@@ -62,6 +62,10 @@ func MapCadenceToIwfSearchAttributes(searchAttributes *shared.SearchAttributes,
 		if err != nil {
 			return nil, err
 		}
+		// when SAs are nil, skip calling mapToIwfSearchAttribute
+		if object == nil {
+			continue
+		}
 		rv, err := mapToIwfSearchAttribute(key, sa.GetValueType(), object, true)
 		if err != nil {
 			return nil, err
@@ -94,6 +98,10 @@ func MapTemporalToIwfSearchAttributes(searchAttributes *common.SearchAttributes,
 		}
 		// TODO we should also call UseNumber here for JSON decoder for Temporal
 		// see https://github.com/temporalio/sdk-go/issues/942
+		// when SAs are nil, skip calling mapToIwfSearchAttribute
+		if object == nil {
+			continue
+		}
 		rv, err := mapToIwfSearchAttribute(key, sa.GetValueType(), object, false)
 		if err != nil {
 			return nil, err
diff --git a/service/interpreter/stateExecutionCounter.go b/service/interpreter/stateExecutionCounter.go
index 3fb43400..1b21ea11 100644
--- a/service/interpreter/stateExecutionCounter.go
+++ b/service/interpreter/stateExecutionCounter.go
@@ -209,6 +209,7 @@ func (e *StateExecutionCounter) refreshIwfExecutingStateIdSearchAttribute() erro
 		})
 		if err != nil {
 			e.provider.GetLogger(e.ctx).Error("error for GetSearchAttributes", err)
+			return err
 		}
 
 		var currentSAsValues []string