From a51aa94c98a3a3eaf8fb5b749ae9a70b08726e92 Mon Sep 17 00:00:00 2001
From: Chris Smalley <csmalley@netflix.com>
Date: Wed, 9 Oct 2019 15:38:15 -0700
Subject: [PATCH] fix(saga-npe): Set sagaContext on DeployDescription object
 where atomic operations operate on TitusDeployHandler (#4084)

---
 .../ops/CloneTitusServerGroupAtomicOperation.groovy  | 12 +++++++++++-
 .../deploy/ops/RunTitusJobAtomicOperation.groovy     | 12 +++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/clouddriver-titus/src/main/groovy/com/netflix/spinnaker/clouddriver/titus/deploy/ops/CloneTitusServerGroupAtomicOperation.groovy b/clouddriver-titus/src/main/groovy/com/netflix/spinnaker/clouddriver/titus/deploy/ops/CloneTitusServerGroupAtomicOperation.groovy
index 6d30ec1c623..1add2f65969 100644
--- a/clouddriver-titus/src/main/groovy/com/netflix/spinnaker/clouddriver/titus/deploy/ops/CloneTitusServerGroupAtomicOperation.groovy
+++ b/clouddriver-titus/src/main/groovy/com/netflix/spinnaker/clouddriver/titus/deploy/ops/CloneTitusServerGroupAtomicOperation.groovy
@@ -20,11 +20,14 @@ import com.netflix.spinnaker.clouddriver.data.task.Task
 import com.netflix.spinnaker.clouddriver.data.task.TaskRepository
 import com.netflix.spinnaker.clouddriver.deploy.DeploymentResult
 import com.netflix.spinnaker.clouddriver.orchestration.AtomicOperation
+import com.netflix.spinnaker.clouddriver.orchestration.SagaContextAware
 import com.netflix.spinnaker.clouddriver.titus.TitusClientProvider
 import com.netflix.spinnaker.clouddriver.titus.deploy.description.TitusDeployDescription
 import com.netflix.spinnaker.clouddriver.titus.deploy.handlers.TitusDeployHandler
 
-class CloneTitusServerGroupAtomicOperation implements AtomicOperation<DeploymentResult> {
+import javax.annotation.Nonnull
+
+class CloneTitusServerGroupAtomicOperation implements AtomicOperation<DeploymentResult>, SagaContextAware {
 
   private static final String PHASE = "CLONE_TITUS_SERVER_GROUP"
 
@@ -50,4 +53,11 @@ class CloneTitusServerGroupAtomicOperation implements AtomicOperation<Deployment
   private static Task getTask() {
     TaskRepository.threadLocalTask.get()
   }
+
+  @Override
+  void setSagaContext(@Nonnull SagaContext sagaContext) {
+    if (description instanceof SagaContextAware) {
+      ((SagaContextAware) description).sagaContext = sagaContext
+    }
+  }
 }
diff --git a/clouddriver-titus/src/main/groovy/com/netflix/spinnaker/clouddriver/titus/deploy/ops/RunTitusJobAtomicOperation.groovy b/clouddriver-titus/src/main/groovy/com/netflix/spinnaker/clouddriver/titus/deploy/ops/RunTitusJobAtomicOperation.groovy
index a7e95ba9476..d575131ac17 100644
--- a/clouddriver-titus/src/main/groovy/com/netflix/spinnaker/clouddriver/titus/deploy/ops/RunTitusJobAtomicOperation.groovy
+++ b/clouddriver-titus/src/main/groovy/com/netflix/spinnaker/clouddriver/titus/deploy/ops/RunTitusJobAtomicOperation.groovy
@@ -20,11 +20,14 @@ import com.netflix.spinnaker.clouddriver.data.task.Task
 import com.netflix.spinnaker.clouddriver.data.task.TaskRepository
 import com.netflix.spinnaker.clouddriver.deploy.DeploymentResult
 import com.netflix.spinnaker.clouddriver.orchestration.AtomicOperation
+import com.netflix.spinnaker.clouddriver.orchestration.SagaContextAware
 import com.netflix.spinnaker.clouddriver.titus.TitusClientProvider
 import com.netflix.spinnaker.clouddriver.titus.deploy.description.TitusDeployDescription
 import com.netflix.spinnaker.clouddriver.titus.deploy.handlers.TitusDeployHandler
 
-class RunTitusJobAtomicOperation implements AtomicOperation<DeploymentResult> {
+import javax.annotation.Nonnull
+
+class RunTitusJobAtomicOperation implements AtomicOperation<DeploymentResult>, SagaContextAware {
 
   private static final String PHASE = "RUN_TITUS_JOB"
 
@@ -51,4 +54,11 @@ class RunTitusJobAtomicOperation implements AtomicOperation<DeploymentResult> {
   private static Task getTask() {
     TaskRepository.threadLocalTask.get()
   }
+
+  @Override
+  void setSagaContext(@Nonnull SagaContext sagaContext) {
+    if (description instanceof SagaContextAware) {
+      ((SagaContextAware) description).sagaContext = sagaContext
+    }
+  }
 }