SERVER-60376 Fix venv setup on container distros

6 files changed, 80 insertions, 49 deletions

diff --git a/etc/evergreen.yml b/etc/evergreen.yml
index 51e6570f8d7..1ea04ccfa20 100644
--- a/etc/evergreen.yml
+++ b/etc/evergreen.yml
@@ -506,17 +506,16 @@ functions:
     params:
       aws_key: ${aws_key}
       aws_secret: ${aws_secret}
-      remote_file: ${project}/${build_variant}/${revision}/mongodb-mongo-venv/${build_id}.tgz
+      remote_file: ${project}/${build_variant}/${revision}/venv/${build_id}.tgz
       bucket: mciuploads
-      # Python venv is not relocatable and should be placed in /data/mci
-      extract_to: ".."
+      extract_to: "."
 
-  "add symlinks back to venv": &add_symlinks_back_to_venv
+  "adjust venv": &adjust_venv
     command: subprocess.exec
     params:
       binary: bash
       args:
-        - "src/evergreen/functions/venv_symlinks_add_back.sh"
+        - "src/evergreen/functions/venv_adjust.sh"
 
   "fetch packages": &fetch_packages
     command: s3.get
@@ -876,7 +875,7 @@ functions:
   - *kill_processes
   - *cleanup_environment
   - *fetch_venv
-  - *add_symlinks_back_to_venv
+  - *adjust_venv
   - *fetch_binaries
   - *extract_binaries
   - *check_binary_version
@@ -1117,7 +1116,7 @@ functions:
     - *kill_processes
     - *cleanup_environment
     - *fetch_venv
-    - *add_symlinks_back_to_venv
+    - *adjust_venv
     - *f_expansions_write
     - *configure_evergreen_api_credentials
 
@@ -2126,22 +2125,23 @@ tasks:
 
     - command: archive.targz_pack
       params:
-        target: "mongodb-mongo-venv.tgz"
-        source_dir: ".."
+        target: "venv.tgz"
+        source_dir: "./"
         include:
-          - "./mongodb-mongo-venv/**"
+          - "./venv/**"
+          - "./venv_readme.txt"
 
     - command: s3.put
       params:
         aws_key: ${aws_key}
         aws_secret: ${aws_secret}
-        local_file: mongodb-mongo-venv.tgz
-        remote_file: ${project}/${build_variant}/${revision}/mongodb-mongo-venv/${build_id}.tgz
+        local_file: venv.tgz
+        remote_file: ${project}/${build_variant}/${revision}/venv/${build_id}.tgz
         bucket: mciuploads
         permissions: private
         visibility: signed
         content_type: application/tar
-        display_name: Python venv (see included README.txt)
+        display_name: Python venv (see included venv_readme.txt)
 
 - name: archive_dist_test_debug
   tags: []
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)"