SERVER-60376 Fix venv setup on container distros

5 files changed, 67 insertions, 36 deletions

diff --git a/evergreen/cleanup_environment.sh b/evergreen/cleanup_environment.sh
index 5755ff93fe9..ec442defb40 100755
--- a/evergreen/cleanup_environment.sh
+++ b/evergreen/cleanup_environment.sh
@@ -1,4 +1,4 @@
 set -o verbose
 
-rm -rf /data/db/* mongo-diskstats* mongo-*.tgz ~/.aws ~/.boto ../mongodb-mongo-venv
+rm -rf /data/db/* mongo-diskstats* mongo-*.tgz ~/.aws ~/.boto venv
 exit 0
diff --git a/evergreen/functions/venv_adjust.sh b/evergreen/functions/venv_adjust.sh
new file mode 100644
index 00000000000..bccdf847f80
--- /dev/null
+++ b/evergreen/functions/venv_adjust.sh
@@ -0,0 +1,33 @@
+DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" > /dev/null 2>&1 && pwd)"
+. "$DIR/../prelude_python.sh"
+
+set -o errexit
+set -o verbose
+
+python_loc=$(which ${python})
+pushd venv
+venv_dir=$(pwd)
+popd
+
+# Update virtual env directory in activate script
+if [ "Windows_NT" = "$OS" ]; then
+  sed -i -e "s:VIRTUAL_ENV=\".*\":VIRTUAL_ENV=\"$venv_dir\":" "$venv_dir/Scripts/activate"
+else
+  sed -i -e "s:VIRTUAL_ENV=\".*\":VIRTUAL_ENV=\"$venv_dir\":" "$venv_dir/bin/activate"
+fi
+
+# Add back python symlinks on linux platforms
+if [ "Windows_NT" = "$OS" ]; then
+  exit 0
+fi
+
+cd "$venv_dir/bin"
+
+rm python python3
+ln -s "$python_loc" python3
+ln -s python3 python
+python3_dot_locs=$(ls python3.*)
+for p in $python3_dot_locs; do
+  rm "$p"
+  ln -s python3 "$p"
+done
diff --git a/evergreen/functions/venv_setup.sh b/evergreen/functions/venv_setup.sh
index 79deed10c43..2b18f07338e 100644
--- a/evergreen/functions/venv_setup.sh
+++ b/evergreen/functions/venv_setup.sh
@@ -6,28 +6,50 @@ evergreen_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" > /dev/null 2>&1 && pwd)/..
 . "$evergreen_dir/prelude_python.sh"
 
 python_loc=$(which ${python})
-venv_dir="${workdir}/../mongodb-mongo-venv"
+venv_dir="${workdir}/venv"
 if [ -d "$venv_dir" ]; then
   exit 0
 fi
 "$python_loc" -m venv "$venv_dir"
 
-pushd $venv_dir
-venv_dir=$(pwd)
-popd
-cat << EOF >> "$venv_dir/README.txt"
+# Adding README file for using this venv locally
+cat << EOF >> venv_readme.txt
 This is an archive of the Python venv generated by this Evergreen build.
 You can use it locally to avoid needing to manually set up the Python environment.
 
-Note that Python venv is not relocatable, you must therefore extract this directory to:
-$venv_dir
+Before activating it you should adjust it for your local environment.
+Run the following commands to do that:
+
+echo "Updating virtual env directory in activate script"
+pushd venv; venv_dir=\$(pwd); popd
+EOF
+if [ "Windows_NT" = "$OS" ]; then
+  cat << EOF >> venv_readme.txt
+sed -i -e "s:VIRTUAL_ENV=\".*\":VIRTUAL_ENV=\"\$venv_dir\":" "\$venv_dir/Scripts/activate"
+EOF
+else
+  cat << EOF >> venv_readme.txt
+sed -i -e "s:VIRTUAL_ENV=\".*\":VIRTUAL_ENV=\"\$venv_dir\":" "\$venv_dir/bin/activate"
+
+echo "Adding back python symlinks"
+pushd venv/bin
+rm python python3
+ln -s "$python_loc" python3
+ln -s python3 python
+python3_dot_locs=\$(ls python3.*)
+for p in \$python3_dot_locs; do
+  rm "\$p"
+  ln -s python3 "\$p"
+done
+popd
 EOF
+fi # End of README file
 
 # venv creates its Scripts/activate file with CLRF endings, which
 # cygwin bash does not like. dos2unix it
 # (See https://bugs.python.org/issue32451)
 if [ "Windows_NT" = "$OS" ]; then
-  dos2unix "${workdir}/../mongodb-mongo-venv/Scripts/activate"
+  dos2unix "${workdir}/venv/Scripts/activate"
 fi
 
 export VIRTUAL_ENV_DISABLE_PROMPT=yes
diff --git a/evergreen/functions/venv_symlinks_add_back.sh b/evergreen/functions/venv_symlinks_add_back.sh
deleted file mode 100644
index 09a507591f5..00000000000
--- a/evergreen/functions/venv_symlinks_add_back.sh
+++ /dev/null
@@ -1,24 +0,0 @@
-DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" > /dev/null 2>&1 && pwd)"
-. "$DIR/../prelude.sh"
-
-# Move to the parent of the Evergreen workdir, which is where the virtualenv is extracted to.
-cd ..
-
-set -o errexit
-set -o verbose
-
-# Ignore Windows since it seems to work.
-if [ "Windows_NT" = "$OS" ]; then
-  exit 0
-fi
-
-cd mongodb-mongo-venv/bin/
-
-rm python
-ln -s "$python" python
-
-pythons=$(ls python3*)
-for p in $pythons; do
-  rm "$p"
-  ln -s python "$p"
-done
diff --git a/evergreen/prelude_venv.sh b/evergreen/prelude_venv.sh
index fd92d10e08a..6b97a8a9cfb 100644
--- a/evergreen/prelude_venv.sh
+++ b/evergreen/prelude_venv.sh
@@ -1,11 +1,11 @@
 function activate_venv {
   # check if virtualenv is set up
-  if [ -d "${workdir}/../mongodb-mongo-venv" ]; then
+  if [ -d "${workdir}/venv" ]; then
     if [ "Windows_NT" = "$OS" ]; then
       # Need to quote the path on Windows to preserve the separator.
-      . "${workdir}/../mongodb-mongo-venv/Scripts/activate" 2> /tmp/activate_error.log
+      . "${workdir}/venv/Scripts/activate" 2> /tmp/activate_error.log
     else
-      . ${workdir}/../mongodb-mongo-venv/bin/activate 2> /tmp/activate_error.log
+      . ${workdir}/venv/bin/activate 2> /tmp/activate_error.log
     fi
     if [ $? -ne 0 ]; then
       echo "Failed to activate virtualenv: $(cat /tmp/activate_error.log)"