From 5afcc25d27b6a7588f3dcf6178fe275eb51cd3f1 Mon Sep 17 00:00:00 2001
From: Cor <jczuurmond@protonmail.com>
Date: Wed, 31 Jul 2024 11:35:06 +0200
Subject: [PATCH] Add upgrade and upgrade eager flags to pip install call
 (#1636)

## Changes
Add upgrade and upgrade eager flags to pip install call for Databricks
labs projects. See [this
documentation](https://pip.pypa.io/en/stable/cli/pip_install/#cmdoption-U)
for more information about the flags.

Resolves #1634

## Tests
- [x] Manually
---
 cmd/labs/project/installer.go      | 4 +++-
 cmd/labs/project/installer_test.go | 8 ++++----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/cmd/labs/project/installer.go b/cmd/labs/project/installer.go
index 39ed9a96..04141596 100644
--- a/cmd/labs/project/installer.go
+++ b/cmd/labs/project/installer.go
@@ -272,8 +272,10 @@ func (i *installer) installPythonDependencies(ctx context.Context, spec string)
 	// - python3 -m ensurepip --default-pip
 	// - curl -o https://bootstrap.pypa.io/get-pip.py | python3
 	var buf bytes.Buffer
+	// Ensure latest version(s) is installed with the `--upgrade` and `--upgrade-strategy eager` flags
+	// https://pip.pypa.io/en/stable/cli/pip_install/#cmdoption-U
 	_, err := process.Background(ctx,
-		[]string{i.virtualEnvPython(ctx), "-m", "pip", "install", spec},
+		[]string{i.virtualEnvPython(ctx), "-m", "pip", "install", "--upgrade", "--upgrade-strategy", "eager", spec},
 		process.WithCombinedOutput(&buf),
 		process.WithDir(libDir))
 	if err != nil {
diff --git a/cmd/labs/project/installer_test.go b/cmd/labs/project/installer_test.go
index 0e049b4c..8754a560 100644
--- a/cmd/labs/project/installer_test.go
+++ b/cmd/labs/project/installer_test.go
@@ -199,7 +199,7 @@ func TestInstallerWorksForReleases(t *testing.T) {
 	stub.WithStdoutFor(`python[\S]+ --version`, "Python 3.10.5")
 	// on Unix, we call `python3`, but on Windows it is `python.exe`
 	stub.WithStderrFor(`python[\S]+ -m venv .*/.databricks/labs/blueprint/state/venv`, "[mock venv create]")
-	stub.WithStderrFor(`python[\S]+ -m pip install .`, "[mock pip install]")
+	stub.WithStderrFor(`python[\S]+ -m pip install --upgrade --upgrade-strategy eager .`, "[mock pip install]")
 	stub.WithStdoutFor(`python[\S]+ install.py`, "setting up important infrastructure")
 
 	// simulate the case of GitHub Actions
@@ -406,7 +406,7 @@ func TestUpgraderWorksForReleases(t *testing.T) {
 	// Install stubs for the python calls we need to ensure were run in the
 	// upgrade process.
 	ctx, stub := process.WithStub(ctx)
-	stub.WithStderrFor(`python[\S]+ -m pip install .`, "[mock pip install]")
+	stub.WithStderrFor(`python[\S]+ -m pip install --upgrade --upgrade-strategy eager .`, "[mock pip install]")
 	stub.WithStdoutFor(`python[\S]+ install.py`, "setting up important infrastructure")
 
 	py, _ := python.DetectExecutable(ctx)
@@ -430,13 +430,13 @@ func TestUpgraderWorksForReleases(t *testing.T) {
 	// Check if the stub was called with the 'python -m pip install' command
 	pi := false
 	for _, call := range stub.Commands() {
-		if strings.HasSuffix(call, "-m pip install .") {
+		if strings.HasSuffix(call, "-m pip install --upgrade --upgrade-strategy eager .") {
 			pi = true
 			break
 		}
 	}
 	if !pi {
-		t.Logf(`Expected stub command 'python[\S]+ -m pip install .' not found`)
+		t.Logf(`Expected stub command 'python[\S]+ -m pip install --upgrade --upgrade-strategy eager .' not found`)
 		t.FailNow()
 	}
 }