From af20271a89e6e98217752c89c6492894f9a66124 Mon Sep 17 00:00:00 2001
From: Quanzheng Long <qlong@indeed.com>
Date: Tue, 7 Feb 2023 10:35:33 -0800
Subject: [PATCH] Fix dataObject/search attributes bugs of overriding each
 others (#45)

---
 integ/persistence_test.go            |  7 +++++++
 integ/persistence_workflow.go        |  4 +++-
 integ/persistence_workflow_state1.go |  7 +++++++
 iwf/persistence_impl.go              | 30 ++++++++++++++--------------
 4 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/integ/persistence_test.go b/integ/persistence_test.go
index 307546f..dd62532 100644
--- a/integ/persistence_test.go
+++ b/integ/persistence_test.go
@@ -45,6 +45,13 @@ func TestPersistenceWorkflow(t *testing.T) {
 	dos[testDataObjectKey].Get(&do)
 	assert.Equal(t, wfId, do.StrValue)
 
+	dos, err = client.GetAllWorkflowDataObjects(context.Background(), wfId, "")
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(dos))
+	var str string
+	dos[testDataObjectKey2].Get(&str)
+	assert.Equal(t, "a string", str)
+
 	sas, err := client.GetWorkflowSearchAttributes(context.Background(), &persistenceWorkflow{}, wfId, "", []string{
 		testSearchAttributeKeyword,
 		testSearchAttributeText,
diff --git a/integ/persistence_workflow.go b/integ/persistence_workflow.go
index 1a3b4c1..e87640c 100644
--- a/integ/persistence_workflow.go
+++ b/integ/persistence_workflow.go
@@ -11,7 +11,8 @@ type persistenceWorkflow struct {
 }
 
 const (
-	testDataObjectKey = "test-data-object"
+	testDataObjectKey  = "test-data-object"
+	testDataObjectKey2 = "test-data-object-2"
 
 	testSearchAttributeInt      = "CustomIntField"
 	testSearchAttributeDatetime = "CustomDatetimeField"
@@ -31,6 +32,7 @@ func (b persistenceWorkflow) GetStates() []iwf.StateDef {
 func (b persistenceWorkflow) GetPersistenceSchema() []iwf.PersistenceFieldDef {
 	return []iwf.PersistenceFieldDef{
 		iwf.DataObjectDef(testDataObjectKey),
+		iwf.DataObjectDef(testDataObjectKey2),
 		iwf.SearchAttributeDef(testSearchAttributeInt, iwfidl.INT),
 		iwf.SearchAttributeDef(testSearchAttributeDatetime, iwfidl.DATETIME),
 		iwf.SearchAttributeDef(testSearchAttributeBool, iwfidl.BOOL),
diff --git a/integ/persistence_workflow_state1.go b/integ/persistence_workflow_state1.go
index 3636aec..9e2e10d 100644
--- a/integ/persistence_workflow_state1.go
+++ b/integ/persistence_workflow_state1.go
@@ -30,6 +30,7 @@ func (b persistenceWorkflowState1) Start(ctx iwf.WorkflowContext, input iwf.Obje
 		panic("this value should be empty because we haven't set it before")
 	}
 	persistence.SetDataObject(testDataObjectKey, do)
+	persistence.SetDataObject(testDataObjectKey2, "a string")
 	persistence.SetSearchAttributeInt(testSearchAttributeInt, 1)
 
 	return iwf.EmptyCommandRequest(), nil
@@ -43,6 +44,12 @@ func (b persistenceWorkflowState1) Decide(ctx iwf.WorkflowContext, input iwf.Obj
 
 	var do ExampleDataObjectModel
 	persistence.GetDataObject(testDataObjectKey, &do)
+	var str string
+	persistence.GetDataObject(testDataObjectKey2, &str)
+	if str != "a string" {
+		panic("testDataObjectKey2 value is incorrect")
+	}
+
 	persistence.SetSearchAttributeDatetime(testSearchAttributeDatetime, do.Datetime)
 	persistence.SetSearchAttributeBool(testSearchAttributeBool, true)
 	return iwf.SingleNextState(persistenceWorkflowState2{}, nil), nil
diff --git a/iwf/persistence_impl.go b/iwf/persistence_impl.go
index 47a71a4..eb4075c 100644
--- a/iwf/persistence_impl.go
+++ b/iwf/persistence_impl.go
@@ -256,55 +256,55 @@ func (p *persistenceImpl) GetToReturn() (
 	searchAttributes []iwfidl.SearchAttribute) {
 	for k, v := range p.dataObjectsToReturn {
 		dataObjectsToReturn = append(dataObjectsToReturn, iwfidl.KeyValue{
-			Key:   &k,
-			Value: &v,
+			Key:   ptr.Any(k),
+			Value: ptr.Any(v),
 		})
 	}
 
 	for k, v := range p.stateLocalToReturn {
 		stateLocalToReturn = append(stateLocalToReturn, iwfidl.KeyValue{
-			Key:   &k,
-			Value: &v,
+			Key:   ptr.Any(k),
+			Value: ptr.Any(v),
 		})
 	}
 
 	for k, v := range p.recordedEvents {
 		recordEvents = append(recordEvents, iwfidl.KeyValue{
-			Key:   &k,
-			Value: &v,
+			Key:   ptr.Any(k),
+			Value: ptr.Any(v),
 		})
 	}
 	for k, sa := range p.saIntToReturn {
 		searchAttributes = append(searchAttributes, iwfidl.SearchAttribute{
-			Key:          &k,
+			Key:          ptr.Any(k),
 			ValueType:    ptr.Any(p.saKeyToType[k]),
-			IntegerValue: &sa,
+			IntegerValue: ptr.Any(sa),
 		})
 	}
 	for k, sa := range p.saStringToReturn {
 		searchAttributes = append(searchAttributes, iwfidl.SearchAttribute{
-			Key:         &k,
+			Key:         ptr.Any(k),
 			ValueType:   ptr.Any(p.saKeyToType[k]),
-			StringValue: &sa,
+			StringValue: ptr.Any(sa),
 		})
 	}
 	for k, sa := range p.saDoubleToReturn {
 		searchAttributes = append(searchAttributes, iwfidl.SearchAttribute{
-			Key:         &k,
+			Key:         ptr.Any(k),
 			ValueType:   ptr.Any(p.saKeyToType[k]),
-			DoubleValue: &sa,
+			DoubleValue: ptr.Any(sa),
 		})
 	}
 	for k, sa := range p.saBoolToReturn {
 		searchAttributes = append(searchAttributes, iwfidl.SearchAttribute{
-			Key:       &k,
+			Key:       ptr.Any(k),
 			ValueType: ptr.Any(p.saKeyToType[k]),
-			BoolValue: &sa,
+			BoolValue: ptr.Any(sa),
 		})
 	}
 	for k, sa := range p.saStrArrToReturn {
 		searchAttributes = append(searchAttributes, iwfidl.SearchAttribute{
-			Key:              &k,
+			Key:              ptr.Any(k),
 			ValueType:        ptr.Any(p.saKeyToType[k]),
 			StringArrayValue: sa,
 		})