Keep using BRANCH for snapshot DRA

Need to use BRANCH for snapshot DRA builds because full version is not
available in artifacts-snapshot API.

Author: jozala

.buildkite/scripts/dra-workflow.sh

@@ -35,10 +35,10 @@ if [[ -n "${VERSION_QUALIFIER:-}" ]]; then
   echo "Version qualifier specified. ES_VERSION=${ES_VERSION}."
 fi
 
-BEATS_BUILD_ID="$(./.ci/scripts/resolve-dra-manifest.sh beats "$ES_VERSION" "$WORKFLOW")"
+BEATS_BUILD_ID="$(./.ci/scripts/resolve-dra-manifest.sh beats "$RM_BRANCH" "$ES_VERSION" "$WORKFLOW")"
 echo "BEATS_BUILD_ID=$BEATS_BUILD_ID"
 
-ML_CPP_BUILD_ID="$(./.ci/scripts/resolve-dra-manifest.sh ml-cpp "$ES_VERSION" "$WORKFLOW")"
+ML_CPP_BUILD_ID="$(./.ci/scripts/resolve-dra-manifest.sh ml-cpp "$RM_BRANCH" "$ES_VERSION" "$WORKFLOW")"
 echo "ML_CPP_BUILD_ID=$ML_CPP_BUILD_ID"
 
 LICENSE_KEY_ARG=""

.ci/scripts/resolve-dra-manifest.sh

@@ -11,14 +11,29 @@ fetch_build() {
 }
 
 ARTIFACT="${ARTIFACT:-$1}"
-ES_VERSION="${ES_VERSION:-$2}"
-WORKFLOW=${WORKFLOW:-$3}
+BRANCH="${BRANCH:-$2}"
+ES_VERSION="${ES_VERSION:-$3}"
+WORKFLOW=${WORKFLOW:-$4}
+
+if [[ "$WORKFLOW" == "staging" ]]
+  LATEST_BUILD=$(fetch_build $WORKFLOW $ARTIFACT $ES_VERSION)
+elif [[ "$WORKFLOW" == "snapshot" ]]
+  LATEST_BUILD=$(fetch_build $WORKFLOW $ARTIFACT $BRANCH)
+else
+  echo "Unknown workflow: $WORKFLOW"
+  exit 1
+fi
 
-LATEST_BUILD=$(fetch_build $WORKFLOW $ARTIFACT $ES_VERSION)
 LATEST_VERSION=$(strip_version $LATEST_BUILD)
 
+# If the latest artifact version doesn't match what we expect, try the corresponding version branch.
+# This can happen when the version of artifact has been bumped on the master branch.
 if [ "$LATEST_VERSION" != "$ES_VERSION" ]; then
-  echo "warning: Latest build for '$ARTIFACT' is version $LATEST_VERSION but expected version $ES_VERSION."
+  echo "Latest build for '$ARTIFACT' is version $LATEST_VERSION but expected version $ES_VERSION." 1>&2
+  NEW_BRANCH=$(echo $ES_VERSION | sed -E "s/([0-9]+\.[0-9]+)\.[0-9]/\1/g")
+
+  echo "Using branch $NEW_BRANCH instead of $BRANCH." 1>&2
+  LATEST_BUILD=$(fetch_build $WORKFLOW $ARTIFACT $NEW_BRANCH)
 fi
 
 echo $LATEST_BUILD