From 35a3ef15f3971e6e373ea913fd91b4c06ae5d859 Mon Sep 17 00:00:00 2001
From: Alex Shtin <alex@shtin.com>
Date: Tue, 6 Sep 2022 16:37:09 -0700
Subject: [PATCH] Rebase

---
 service/frontend/workflow_handler.go | 85 ++++++++++++++++++----------
 1 file changed, 55 insertions(+), 30 deletions(-)

diff --git a/service/frontend/workflow_handler.go b/service/frontend/workflow_handler.go
index eeac49bc1ab6..e8e76bb64b80 100644
--- a/service/frontend/workflow_handler.go
+++ b/service/frontend/workflow_handler.go
@@ -2908,12 +2908,12 @@ func (wh *WorkflowHandler) CreateSchedule(ctx context.Context, request *workflow
 		return nil, err
 	}
 
-	request, err = wh.mapCreateScheduleRequestSearchAttributes(request, namespaceName)
-	if err != nil {
+	if err = wh.validateStartWorkflowArgsForSchedule(namespaceName, request.GetSchedule().GetAction().GetStartWorkflow()); err != nil {
 		return nil, err
 	}
 
-	if err = wh.validateStartWorkflowArgsForSchedule(namespaceName, request.GetSchedule().GetAction().GetStartWorkflow()); err != nil {
+	request, err = wh.mapCreateScheduleRequestSearchAttributes(request, namespaceName)
+	if err != nil {
 		return nil, err
 	}
 
@@ -3005,15 +3005,8 @@ func (wh *WorkflowHandler) validateStartWorkflowArgsForSchedule(
 		return errIDReusePolicyNotAllowed
 	}
 
-	err = wh.validateSearchAttributes(startWorkflow.GetSearchAttributes(), namespaceName)
-	if err != nil {
-		return nil, err
-	}
-
-	// map search attributes to aliases here, since we don't go through the frontend when starting later
-	request, err = wh.substituteSearchAttributesAliasesInCreateScheduleRequest(request, namespaceName)
-	if err != nil {
-		return nil, err
+	if err := wh.validateSearchAttributes(startWorkflow.GetSearchAttributes(), namespaceName); err != nil {
+		return err
 	}
 
 	return nil
@@ -3110,9 +3103,13 @@ func (wh *WorkflowHandler) DescribeSchedule(ctx context.Context, request *workfl
 				return serviceerror.NewUnavailable(fmt.Sprintf(errUnableToGetSearchAttributesMessage, err))
 			}
 			searchattribute.ApplyTypeMap(sa, saTypeMap)
-			if err = searchattribute.ApplyAliases(wh.saMapper, sa, request.Namespace); err != nil {
+			mappedSearchAttributes, err := searchattribute.ApplyAliases(wh.saMapper, sa, request.Namespace)
+			if err != nil {
 				return err
 			}
+			if mappedSearchAttributes != nil {
+				response.Schedule.Action.GetStartWorkflow().SearchAttributes = mappedSearchAttributes
+			}
 		}
 
 		// for all running workflows started by the schedule, we should check that they're
@@ -3230,6 +3227,11 @@ func (wh *WorkflowHandler) UpdateSchedule(ctx context.Context, request *workflow
 		return nil, err
 	}
 
+	request, err = wh.mapUpdateScheduleRequestStartWorkflowSearchAttributes(request, namespaceName)
+	if err != nil {
+		return nil, err
+	}
+
 	input := &schedspb.FullUpdateRequest{
 		Schedule: request.Schedule,
 	}
@@ -4753,36 +4755,59 @@ func (wh *WorkflowHandler) mapCreateScheduleRequestSearchAttributes(request *wor
 	if err != nil {
 		return nil, err
 	}
-	if mappedSearchAttributes == nil {
+
+	startWorkflow := request.GetSchedule().GetAction().GetStartWorkflow()
+	mappedStartWorkflowSearchAttributes, err := searchattribute.SubstituteAliases(wh.saMapper, startWorkflow.GetSearchAttributes(), namespaceName.String())
+	if err != nil {
+		return nil, err
+	}
+
+	if mappedSearchAttributes == nil && mappedStartWorkflowSearchAttributes == nil {
 		return request, nil
 	}
 
 	// Shallow copy request and replace SearchAttributes fields only.
 	newRequest := *request
-	newRequest.SearchAttributes = mappedSearchAttributes
-	return &newRequest, nil
-}
 
-func (wh *WorkflowHandler) mapCreateScheduleRequestStartWorkflowSearchAttributes(request *workflowservice.CreateScheduleRequest, namespaceName namespace.Name) (*workflowservice.CreateScheduleRequest, error) {
-	if startWorkflow := request.GetSchedule().GetAction().GetStartWorkflow(); startWorkflow != nil {
-		mappedSearchAttributes, err := searchattribute.SubstituteAliases(wh.saMapper, startWorkflow.GetSearchAttributes(), namespaceName.String())
-		if err != nil {
-			return nil, err
-		}
-		if mappedSearchAttributes == nil {
-			return request, nil
-		}
+	if mappedSearchAttributes != nil {
+		newRequest.SearchAttributes = mappedSearchAttributes
+	}
+
+	if mappedStartWorkflowSearchAttributes != nil && startWorkflow != nil {
 		newStartWorkflow := *startWorkflow
-		newStartWorkflow.SearchAttributes = mappedSearchAttributes
+		newStartWorkflow.SearchAttributes = mappedStartWorkflowSearchAttributes
 		newSchedule := *request.GetSchedule()
 		newSchedule.Action = &schedpb.ScheduleAction{
 			Action: &schedpb.ScheduleAction_StartWorkflow{
 				StartWorkflow: &newStartWorkflow,
 			}}
-		newRequest := *request
 		newRequest.Schedule = &newSchedule
-		return &newRequest, nil
 	}
 
-	return request, nil
+	return &newRequest, nil
+}
+
+func (wh *WorkflowHandler) mapUpdateScheduleRequestStartWorkflowSearchAttributes(request *workflowservice.UpdateScheduleRequest, namespaceName namespace.Name) (*workflowservice.UpdateScheduleRequest, error) {
+	startWorkflow := request.GetSchedule().GetAction().GetStartWorkflow()
+	if startWorkflow == nil {
+		return request, nil
+	}
+
+	mappedSearchAttributes, err := searchattribute.SubstituteAliases(wh.saMapper, startWorkflow.GetSearchAttributes(), namespaceName.String())
+	if err != nil {
+		return nil, err
+	}
+	if mappedSearchAttributes == nil {
+		return request, nil
+	}
+	newStartWorkflow := *startWorkflow
+	newStartWorkflow.SearchAttributes = mappedSearchAttributes
+	newSchedule := *request.GetSchedule()
+	newSchedule.Action = &schedpb.ScheduleAction{
+		Action: &schedpb.ScheduleAction_StartWorkflow{
+			StartWorkflow: &newStartWorkflow,
+		}}
+	newRequest := *request
+	newRequest.Schedule = &newSchedule
+	return &newRequest, nil
 }