send stage type from TAS to the host and setup the environment variable

deploy-agent/deployd/__init__.py

@@ -27,4 +27,4 @@
 # 2: puppet applied successfully with changes
 PUPPET_SUCCESS_EXIT_CODES = [0, 2]
 
-__version__ = '1.2.53'
+__version__ = '1.2.54'

deploy-agent/deployd/common/config.py

@@ -104,6 +104,8 @@ def update_variables(self, deploy_status):
             self._environ['ENV_NAME'] = deploy_status.report.envName
         if deploy_status.report.stageName:
             self._environ['STAGE_NAME'] = deploy_status.report.stageName
+        if deploy_status.report.stageType:
+            self._environ['COMPUTE_STAGE_TYPE'] = deploy_status.report.stageType
         if deploy_status.first_deploy:
             self._environ['FIRST_DEPLOY'] = str(deploy_status.first_deploy)
         if deploy_status.is_docker:

deploy-agent/deployd/common/types.py

@@ -141,6 +141,7 @@ def update_by_response(self, response):
         self.report.envName = deploy_goal.envName
         self.report.status = AgentStatus.UNKNOWN
         self.report.stageName = deploy_goal.stageName
+        self.report.stageType = deploy_goal.stageType
         self.first_deploy = deploy_goal.firstDeploy
         self.is_docker = deploy_goal.isDocker
 

deploy-agent/deployd/types/deploy_goal.py

@@ -21,6 +21,7 @@ def __init__(self, jsonValue=None):
         self.envId = None
         self.envName = None
         self.stageName = None
+        self.stageType = None
         self.deployStage = None
         self.build = None
         self.deployAlias = None
@@ -34,6 +35,7 @@ def __init__(self, jsonValue=None):
             self.envId = jsonValue.get('envId')
             self.envName = jsonValue.get('envName')
             self.stageName = jsonValue.get('stageName')
+            self.stageType = jsonValue.get('stageType')
             # TODO: Only used for migration, should remove later
             if isinstance(jsonValue.get('deployStage'), int):
                 self.deployStage = DeployStage._VALUES_TO_NAMES[jsonValue.get('deployStage')]
@@ -54,6 +56,7 @@ def __key(self):
                 self.envId,
                 self.envName,
                 self.stageName,
+                self.stageType,
                 self.deployStage,
                 self.build,
                 self.deployAlias,
@@ -76,10 +79,10 @@ def __ne__(self, other):
                     and self.__key() == other.__key())
 
     def __str__(self):
-        return "DeployGoal(deployId={}, envId={}, envName={}, stageName={}, " \
+        return "DeployGoal(deployId={}, envId={}, envName={}, stageName={}, stageType={}, " \
                "deployStage={}, build={}, deployAlias={}, agentConfig={}," \
                "scriptVariables={}, firstDeploy={}, isDocker={})".format(self.deployId, self.envId, self.envName,
-                                                            self.stageName, self.deployStage,
+                                                            self.stageName, self.stageType, self.deployStage,
                                                             self.build, self.deployAlias,
                                                             self.config, self.scriptVariables,
                                                             self.firstDeploy, self.isDocker)

deploy-agent/deployd/types/ping_report.py

@@ -23,6 +23,7 @@ def __init__(self, jsonValue=None):
         self.envId = None
         self.envName = None
         self.stageName = None
+        self.stageType = None
         self.deployStage = None
         self.status = None
         self.errorCode = 0
@@ -50,6 +51,7 @@ def __init__(self, jsonValue=None):
 
             self.envName = jsonValue.get('envName')
             self.stageName = jsonValue.get('stageName')
+            self.stageType = jsonValue.get('stageType')
             self.errorCode = jsonValue.get('errorCode')
             self.errorMessage = jsonValue.get('errorMessage')
             self.failCount = jsonValue.get('failCount')

deploy-agent/tests/unit/deploy/common/test_config.py

@@ -46,6 +46,7 @@ def test_get_target(self):
         deploy_goal['deployId'] = '123'
         deploy_goal['stageName'] = 'beta'
         deploy_goal['envName'] = 'pinboard'
+        deploy_goal['stageType'] = 'DEFAULT'
         deploy_goal['deployStage'] = DeployStage.SERVING_BUILD
         ping_response = {'deployGoal': deploy_goal, 'opCode': OpCode.NOOP}
 
@@ -54,6 +55,7 @@ def test_get_target(self):
         self.assertEqual(os.environ['DEPLOY_ID'], '123')
         self.assertEqual(os.environ['ENV_NAME'], 'pinboard')
         self.assertEqual(os.environ['STAGE_NAME'], 'beta')
+        self.assertEqual(os.environ['COMPUTE_STAGE_TYPE'], 'DEFAULT')
         self.assertEqual(self.config.get_target(), '/tmp/pinboard')
 
 

deploy-service/common/src/main/java/com/pinterest/deployservice/bean/DeployGoalBean.java

@@ -25,6 +25,7 @@ public class DeployGoalBean {
     private String envId;
     private String envName;
     private String stageName;
+    private EnvType stageType;
     private DeployStage deployStage;
     private BuildBean build;
     private String deployAlias;
@@ -65,6 +66,14 @@ public void setStageName(String stageName) {
         this.stageName = stageName;
     }
 
+    public EnvType getStageType() {
+        return stageType;
+    }
+
+    public void setStageType(EnvType stageType) {
+        this.stageType = stageType;
+    }
+
     public DeployStage getDeployStage() {
         return deployStage;
     }

deploy-service/common/src/main/java/com/pinterest/deployservice/handler/PingHandler.java

@@ -760,6 +760,8 @@ public PingResponseBean generateInstallResponse(GoalAnalyst.InstallCandidate ins
         goal.setEnvId(envBean.getEnv_id());
         goal.setEnvName(envBean.getEnv_name());
         goal.setStageName(envBean.getStage_name());
+        LOG.debug("stage type: {}", envBean.getStage_type());
+        goal.setStageType(envBean.getStage_type());
         goal.setIsDocker(envBean.getIs_docker());
 
         // TODO optimize the next stage here based on deploy ( some deploy does not have all the stages )
@@ -835,6 +837,7 @@ PingResponseBean generateDeleteResponse(GoalAnalyst.UninstallCandidate candidate
         goal.setDeployStage(candidate.report.getDeployStage());
         goal.setEnvName(candidate.environ.getEnv_name());
         goal.setStageName(candidate.environ.getStage_name());
+        goal.setStageType(candidate.environ.getStage_type());
         goal.setFirstDeploy(false);
         response.setDeployGoal(goal);
         return response;