installer: use venv and improve error handling (#4099)

* installer: improve error handling and logging

Co-authored-by: Bjorn Neergaard <bjorn@neersighted.com>

* installer: default to using in-built venv module

This change refactors virtual environment creation and introduces a new
wrapper to handle command execution etc. If venv module is not
available, in cases like that of the ubuntu distribution where the
module is not part of the default python install, the virtualenv
package is used by fetching the appropriate zipapp for the package for
the python runtime.

Resolves: #4089
Resolves: #4093

* ci/installer: upload failure logs

* installer: preserve existing environment on failure

This change ensures that if creation of a new environment fails, any
previously existing environment is restored.

Co-authored-by: Bjorn Neergaard <bjorn@neersighted.com>

.github/workflows/installer.yml

@@ -49,6 +49,13 @@ jobs:
         shell: bash
         run: python install-poetry.py -y
 
+      - name: Upload Failure Log
+        uses: actions/upload-artifact@v2
+        if: failure()
+        with:
+          name: poetry-installer-error.log
+          path: poetry-installer-error-*.log
+
       - name: Verify Installation
         shell: bash
         run: |

install-poetry.py

@@ -3,14 +3,15 @@
 
 It does, in order:
 
-  - Downloads the virtualenv package to a temporary directory and add it to sys.path.
-  - Creates a virtual environment in the correct OS data dir which will be
+  - Creates a virtual environment using venv (or virtualenv zipapp) in the correct OS data dir which will be
       - `%APPDATA%\\pypoetry` on Windows
       -  ~/Library/Application Support/pypoetry on MacOS
       - `${XDG_DATA_HOME}/pypoetry` (or `~/.local/share/pypoetry` if it's not set) on UNIX systems
       - In `${POETRY_HOME}` if it's set.
   - Installs the latest or given version of Poetry inside this virtual environment.
   - Installs a `poetry` script in the Python user directory (or `${POETRY_HOME/bin}` if `POETRY_HOME` is set).
+  - On failure, the error log is written to poetry-installer-error-*.log and any previously existing environment
+    is restored.
 """
 
 import argparse
@@ -219,21 +220,6 @@ def _get_win_folder_with_ctypes(csidl_name):
         _get_win_folder = _get_win_folder_from_registry
 
 
-@contextmanager
-def temporary_directory(*args, **kwargs):
-    try:
-        from tempfile import TemporaryDirectory
-    except ImportError:
-        name = tempfile.mkdtemp(*args, **kwargs)
-
-        yield name
-
-        shutil.rmtree(name)
-    else:
-        with TemporaryDirectory(*args, **kwargs) as name:
-            yield name
-
-
 PRE_MESSAGE = """# Welcome to {poetry}!
 
 This will download and install the latest version of {poetry},
@@ -277,6 +263,83 @@ def temporary_directory(*args, **kwargs):
 POST_MESSAGE_CONFIGURE_WINDOWS = """"""
 
 
+class PoetryInstallationError(RuntimeError):
+    def __init__(self, return_code: int = 0, log: Optional[str] = None):
+        super(PoetryInstallationError, self).__init__()
+        self.return_code = return_code
+        self.log = log
+
+
+class VirtualEnvironment:
+    def __init__(self, path: Path) -> None:
+        self._path = path
+        # str is required for compatibility with subprocess run on CPython <= 3.7 on Windows
+        self._python = str(
+            self._path.joinpath("Scripts/python.exe" if WINDOWS else "bin/python")
+        )
+
+    @property
+    def path(self):
+        return self._path
+
+    @classmethod
+    def make(cls, target: Path) -> "VirtualEnvironment":
+        try:
+            import venv
+
+            builder = venv.EnvBuilder(clear=True, with_pip=True, symlinks=False)
+            builder.ensure_directories(target)
+            builder.create(target)
+        except ImportError:
+            # fallback to using virtualenv package if venv is not available, eg: ubuntu
+            python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+            virtualenv_bootstrap_url = (
+                f"https://bootstrap.pypa.io/virtualenv/{python_version}/virtualenv.pyz"
+            )
+
+            with tempfile.TemporaryDirectory(prefix="poetry-installer") as temp_dir:
+                virtualenv_pyz = Path(temp_dir) / "virtualenv.pyz"
+                request = Request(
+                    virtualenv_bootstrap_url, headers={"User-Agent": "Python Poetry"}
+                )
+                virtualenv_pyz.write_bytes(urlopen(request).read())
+                cls.run(
+                    sys.executable, virtualenv_pyz, "--clear", "--always-copy", target
+                )
+
+        # We add a special file so that Poetry can detect
+        # its own virtual environment
+        target.joinpath("poetry_env").touch()
+
+        env = cls(target)
+
+        # we do this here to ensure that outdated system default pip does not trigger older bugs
+        env.pip("install", "--disable-pip-version-check", "--upgrade", "pip")
+
+        return env
+
+    @staticmethod
+    def run(*args, **kwargs) -> subprocess.CompletedProcess:
+        completed_process = subprocess.run(
+            args,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            **kwargs,
+        )
+        if completed_process.returncode != 0:
+            raise PoetryInstallationError(
+                return_code=completed_process.returncode,
+                log=completed_process.stdout.decode(),
+            )
+        return completed_process
+
+    def python(self, *args, **kwargs) -> subprocess.CompletedProcess:
+        return self.run(self._python, *args, **kwargs)
+
+    def pip(self, *args, **kwargs) -> subprocess.CompletedProcess:
+        return self.python("-m", "pip", "--isolated", *args, **kwargs)
+
+
 class Cursor:
     def __init__(self) -> None:
         self._output = sys.stdout
@@ -439,12 +502,10 @@ def _is_self_upgrade_supported(x):
         try:
             self.install(version)
         except subprocess.CalledProcessError as e:
-            print(
-                colorize("error", f"\nAn error has occurred: {e}\n{e.stdout.decode()}")
+            raise PoetryInstallationError(
+                return_code=e.returncode, log=e.output.decode()
             )
 
-            return e.returncode
-
         self._write("")
         self.display_post_message(version)
 
@@ -460,21 +521,13 @@ def install(self, version, upgrade=False):
             )
         )
 
-        env_path = self.make_env(version)
-        self.install_poetry(version, env_path)
-        self.make_bin(version)
+        with self.make_env(version) as env:
+            self.install_poetry(version, env)
+            self.make_bin(version, env)
+            self._data_dir.joinpath("VERSION").write_text(version)
+            self._install_comment(version, "Done")
 
-        self._overwrite(
-            "Installing {} ({}): {}".format(
-                colorize("info", "Poetry"),
-                colorize("b", version),
-                colorize("success", "Done"),
-            )
-        )
-
-        self._data_dir.joinpath("VERSION").write_text(version)
-
-        return 0
+            return 0
 
     def uninstall(self) -> int:
         if not self._data_dir.exists():
@@ -504,81 +557,70 @@ def uninstall(self) -> int:
 
         return 0
 
-    def make_env(self, version: str) -> Path:
+    def _install_comment(self, version: str, message: str):
         self._overwrite(
             "Installing {} ({}): {}".format(
                 colorize("info", "Poetry"),
                 colorize("b", version),
-                colorize("comment", "Creating environment"),
+                colorize("comment", message),
             )
         )
 
+    @contextmanager
+    def make_env(self, version: str) -> VirtualEnvironment:
         env_path = self._data_dir.joinpath("venv")
+        env_path_saved = env_path.with_suffix(".save")
 
-        with temporary_directory() as tmp_dir:
-            subprocess.run(
-                [sys.executable, "-m", "pip", "install", "virtualenv", "-t", tmp_dir],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                check=True,
-            )
-
-            sys.path.insert(0, tmp_dir)
-
-            import virtualenv
+        if env_path.exists():
+            self._install_comment(version, "Saving existing environment")
+            if env_path_saved.exists():
+                shutil.rmtree(env_path_saved)
+            shutil.move(env_path, env_path_saved)
 
-            virtualenv.cli_run([str(env_path), "--clear"])
-
-        # We add a special file so that Poetry can detect
-        # its own virtual environment
-        env_path.joinpath("poetry_env").touch()
+        try:
+            self._install_comment(version, "Creating environment")
+            yield VirtualEnvironment.make(env_path)
+        except Exception as e:  # noqa
+            if env_path.exists():
+                self._install_comment(
+                    version, "An error occurred. Removing partial environment."
+                )
+                shutil.rmtree(env_path)
 
-        return env_path
+            if env_path_saved.exists():
+                self._install_comment(
+                    version, "Restoring previously saved environment."
+                )
+                shutil.move(env_path_saved, env_path)
 
-    def make_bin(self, version: str) -> None:
-        self._overwrite(
-            "Installing {} ({}): {}".format(
-                colorize("info", "Poetry"),
-                colorize("b", version),
-                colorize("comment", "Creating script"),
-            )
-        )
+            raise e
+        else:
+            if env_path_saved.exists():
+                shutil.rmtree(env_path_saved, ignore_errors=True)
 
+    def make_bin(self, version: str, env: VirtualEnvironment) -> None:
+        self._install_comment(version, "Creating script")
         self._bin_dir.mkdir(parents=True, exist_ok=True)
 
         script = "poetry"
-        target_script = "venv/bin/poetry"
+        script_bin = "bin"
         if WINDOWS:
             script = "poetry.exe"
-            target_script = "venv/Scripts/poetry.exe"
+            script_bin = "Scripts"
+        target_script = env.path.joinpath(script_bin, script)
 
         if self._bin_dir.joinpath(script).exists():
             self._bin_dir.joinpath(script).unlink()
 
         try:
-            self._bin_dir.joinpath(script).symlink_to(
-                self._data_dir.joinpath(target_script)
-            )
+            self._bin_dir.joinpath(script).symlink_to(target_script)
         except OSError:
             # This can happen if the user
             # does not have the correct permission on Windows
-            shutil.copy(
-                self._data_dir.joinpath(target_script), self._bin_dir.joinpath(script)
-            )
+            shutil.copy(target_script, self._bin_dir.joinpath(script))
 
-    def install_poetry(self, version: str, env_path: Path) -> None:
-        self._overwrite(
-            "Installing {} ({}): {}".format(
-                colorize("info", "Poetry"),
-                colorize("b", version),
-                colorize("comment", "Installing Poetry"),
-            )
-        )
-
-        if WINDOWS:
-            python = env_path.joinpath("Scripts/python.exe")
-        else:
-            python = env_path.joinpath("bin/python")
+    def install_poetry(self, version: str, env: VirtualEnvironment) -> None:
+        self._install_comment(version, "Installing Poetry")
 
         if self._git:
             specification = "git+" + version
@@ -587,12 +629,7 @@ def install_poetry(self, version: str, env_path: Path) -> None:
         else:
             specification = f"poetry=={version}"
 
-        subprocess.run(
-            [str(python), "-m", "pip", "install", specification],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            check=True,
-        )
+        env.pip("install", specification)
 
     def display_pre_message(self) -> None:
         kwargs = {
@@ -831,7 +868,25 @@ def main():
     if args.uninstall or string_to_bool(os.getenv("POETRY_UNINSTALL", "0")):
         return installer.uninstall()
 
-    return installer.run()
+    try:
+        return installer.run()
+    except PoetryInstallationError as e:
+        installer._write(colorize("error", "Poetry installation failed."))  # noqa
+
+        if e.log is not None:
+            import traceback
+
+            _, path = tempfile.mkstemp(
+                suffix=".log",
+                prefix="poetry-installer-error-",
+                dir=str(Path.cwd()),
+                text=True,
+            )
+            installer._write(colorize("error", f"See {path} for error logs."))  # noqa
+            text = f"{e.log}\nTraceback:\n\n{''.join(traceback.format_tb(e.__traceback__))}"
+            Path(path).write_text(text)
+
+        return e.return_code
 
 
 if __name__ == "__main__":