From 4c878b81203e29eea312b5774868b83c1b931466 Mon Sep 17 00:00:00 2001
From: Mike Sarahan <msarahan@gmail.com>
Date: Tue, 16 Apr 2024 17:17:16 -0500
Subject: [PATCH] add rapids-get-pr-wheel-artifact for quickly downloading cpp
 wheels (#101)

---
 tools/rapids-download-wheels-from-s3 | 10 +++++++-
 tools/rapids-get-pr-wheel-artifact   | 36 ++++++++++++++++++++++++++++
 tools/rapids-package-name            |  4 ++++
 tools/rapids-upload-wheels-to-s3     | 12 ++++++++--
 4 files changed, 59 insertions(+), 3 deletions(-)
 create mode 100755 tools/rapids-get-pr-wheel-artifact

diff --git a/tools/rapids-download-wheels-from-s3 b/tools/rapids-download-wheels-from-s3
index a54556b..ee60297 100755
--- a/tools/rapids-download-wheels-from-s3
+++ b/tools/rapids-download-wheels-from-s3
@@ -6,6 +6,14 @@ set -eo pipefail
 
 source rapids-prompt-local-repo-config
 
-pkg_name="$(rapids-package-name wheel_python)"
+# For legacy reasons, allow this script to be run without the pkg_type being the first arg
+pkg_name="$(rapids-package-name "wheel_python")"
+
+pkg_type="$1"
+if [ "${pkg_type}" = "cpp" ] || [ "${pkg_type}" = "python" ]; then
+    # remove pkg_type from args because we handle it in this script
+    shift;
+    pkg_name="$(rapids-package-name "wheel_${pkg_type}")"
+fi
 
 rapids-download-from-s3 "${pkg_name}" "$@"
diff --git a/tools/rapids-get-pr-wheel-artifact b/tools/rapids-get-pr-wheel-artifact
new file mode 100755
index 0000000..548f9d4
--- /dev/null
+++ b/tools/rapids-get-pr-wheel-artifact
@@ -0,0 +1,36 @@
+#!/bin/bash
+# Echo path to an artifact for a specific PR. Finds and uses the latest commit on the PR.
+#
+# Positional Arguments:
+#   1) repo name
+#   2) PR number
+#   3) "cpp" or "python", to get the artifact for the C++ or Python build, respectively
+#   4) [optional] commit hash, to get the artifact for a specific commit
+#
+# Example Usage:
+#   rapids-get-pr-conda-artifact rmm 1095 cpp
+set -euo pipefail
+
+repo="$1"
+pr="$2"
+
+pkg_type="$3"
+case "${pkg_type}" in
+  cpp)
+    artifact_name=$(RAPIDS_REPOSITORY=$repo rapids-package-name wheel_cpp)
+    ;;
+  python)
+    artifact_name=$(RAPIDS_REPOSITORY=$repo rapids-package-name wheel_python)
+    ;;
+  *)
+    echo "Error: 3rd argument must be 'cpp' or 'python'"
+    exit 1
+    ;;
+esac
+
+commit="${4:-}"
+if [[ -z "${commit}" ]]; then
+    commit=$(git ls-remote https://github.com/rapidsai/"${repo}".git refs/heads/pull-request/"${pr}" | cut -c1-7)
+fi
+
+rapids-get-artifact "ci/${repo}/pull-request/${pr}/${commit}/${artifact_name}"
diff --git a/tools/rapids-package-name b/tools/rapids-package-name
index 9899e0c..eaadc5c 100755
--- a/tools/rapids-package-name
+++ b/tools/rapids-package-name
@@ -29,6 +29,10 @@ case "${pkg_type}" in
     append_pyver=1
     append_arch=1
     ;;
+  wheel_cpp)
+    append_wheelname=1
+    append_arch=1
+    ;;
   wheel_python)
     append_wheelname=1
     # Pure wheels do not need a pyver or arch
diff --git a/tools/rapids-upload-wheels-to-s3 b/tools/rapids-upload-wheels-to-s3
index 0dfe33a..9b35b39 100755
--- a/tools/rapids-upload-wheels-to-s3
+++ b/tools/rapids-upload-wheels-to-s3
@@ -4,12 +4,20 @@
 #   1) wheel path to tar up
 set -e
 
+# For legacy reasons, allow this script to be run without the pkg_type being the first arg
+pkg_name="$(rapids-package-name "wheel_python")"
+
+pkg_type="$1"
+if [ "${pkg_type}" = "cpp" ] || [ "${pkg_type}" = "python" ]; then
+    # remove pkg_type from args because we handle it in this script
+    shift;
+    pkg_name="$(rapids-package-name "wheel_${pkg_type}")"
+fi
+
 if [ "${CI:-false}" = "false" ]; then
   rapids-echo-stderr "Packages from local builds cannot be uploaded to S3."
   rapids-echo-stderr "Open a PR to have successful builds uploaded."
   exit 0
 fi
 
-pkg_name="$(rapids-package-name wheel_python)"
-
 rapids-upload-to-s3 "${pkg_name}" "$@"