Fix processing short pip flags in environment dependencies (#3708)

## Changes
Fix processing short pip flags in environment dependencies

## Why
Fixes #3674

## Tests
Added unit test case

<!-- If your PR needs to be included in the release notes for next
release,
add a separate entry in NEXT_CHANGELOG.md as part of your PR. -->

Author: andrewnester

NEXT_CHANGELOG.md

@@ -9,5 +9,7 @@
 ### Dependency updates
 
 ### Bundles
+* Fix processing short pip flags in environment dependencies ([#3708](https://github.com/databricks/cli/pull/3708))
+* Add support for referencing local files in -e pip flag for environment dependencies ([#3708](https://github.com/databricks/cli/pull/3708))
 
 ### API Changes

acceptance/bundle/environments/dependencies/databricks.yml

@@ -21,6 +21,9 @@ resources:
             client: "1"
             dependencies:
               - "-r ./requirements.txt"
+              - "-e ./file.py"
+              - "-i http://myindexurl.com"
+              - "--index-url http://myindexurl.com"
               - "test_package"
               - "test_package==2.0.1"
               - "test_package>=2.0.1"

acceptance/bundle/environments/dependencies/file.py

@@ -23,6 +23,9 @@ Deployment complete!
       "client": "1",
       "dependencies": [
         "-r /Workspace/Users/[USERNAME]/.bundle/dependencies/default/files/requirements.txt",
+        "-e /Workspace/Users/[USERNAME]/.bundle/dependencies/default/files/file.py",
+        "-i http://myindexurl.com",
+        "--index-url http://myindexurl.com",
         "test_package",
         "test_package==2.0.1",
         "test_package>=2.0.1",
@@ -79,6 +82,9 @@ Deployment complete!
       "client": "1",
       "dependencies": [
         "-r /Workspace/Users/[USERNAME]/.bundle/dependencies/default/files/requirements.txt",
+        "-e /Workspace/Users/[USERNAME]/.bundle/dependencies/default/files/file.py",
+        "-i http://myindexurl.com",
+        "--index-url http://myindexurl.com",
         "test_package",
         "test_package==2.0.1",
         "test_package>=2.0.1",

acceptance/bundle/environments/dependencies/output.txt

@@ -11,6 +11,7 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config/mutator/paths"
+	"github.com/databricks/cli/bundle/libraries"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/dyn"
 )
@@ -85,18 +86,20 @@ func collectGitSourcePaths(b *bundle.Bundle) []dyn.Path {
 }
 
 func normalizePath(path string, location dyn.Location, bundleRootPath string) (string, error) {
-	// Handle requirements file paths with -r flag
-	reqPath, ok := strings.CutPrefix(path, "-r ")
-	if ok {
-		// Normalize the path part
-		reqPath = strings.TrimSpace(reqPath)
-		normalizedPath, err := normalizePath(reqPath, location, bundleRootPath)
-		if err != nil {
-			return "", err
-		}
+	// Handle local file paths used inside pip flags
+	for _, flag := range libraries.PipFlagsWithLocalPaths {
+		reqPath, ok := strings.CutPrefix(path, flag+" ")
+		if ok {
+			// Normalize the path part
+			reqPath = strings.TrimSpace(reqPath)
+			normalizedPath, err := normalizePath(reqPath, location, bundleRootPath)
+			if err != nil {
+				return "", err
+			}
 
-		// Reconstruct the path with -r flag
-		return "-r " + normalizedPath, nil
+			// Reconstruct the path with -r flag
+			return flag + " " + normalizedPath, nil
+		}
 	}
 
 	pathAsUrl, err := url.Parse(path)

bundle/config/mutator/normalize_paths.go

@@ -78,6 +78,14 @@ func TestNormalizePath_requirementsFile(t *testing.T) {
 	assert.Equal(t, "-r requirements.txt", value)
 }
 
+func TestNormalizePath_environmentDependency(t *testing.T) {
+	tmpDir := t.TempDir()
+	location := dyn.Location{File: filepath.Join(tmpDir, "resources", "job_1.yml")}
+	value, err := normalizePath("-e ../file.py", location, tmpDir)
+	assert.NoError(t, err)
+	assert.Equal(t, "-e file.py", value)
+}
+
 func TestLocationDirectory(t *testing.T) {
 	loc := dyn.Location{File: "file", Line: 1, Column: 2}
 	dir, err := locationDirectory(loc)

bundle/config/mutator/normalize_paths_test.go

@@ -62,9 +62,9 @@ func jobLibrariesRewritePatterns() []jobRewritePattern {
 				dyn.Key("dependencies"),
 				dyn.AnyIndex(),
 			),
-			TranslateModeEnvironmentRequirements,
+			TranslateModeEnvironmentPipFlag,
 			func(s string) bool {
-				_, ok := libraries.IsLocalRequirementsFile(s)
+				_, _, ok := libraries.IsLocalPathInPipFlag(s)
 				return !ok
 			},
 		},

bundle/config/mutator/paths/job_libraries_paths_visitor.go

@@ -30,6 +30,6 @@ const (
 	// This allows for disambiguating between paths and PyPI package names.
 	TranslateModeLocalRelativeWithPrefix
 
-	// TranslateModeEnvironmentRequirements translates a local requirements file path to be absolute.
-	TranslateModeEnvironmentRequirements
+	// TranslateModeEnvironmentPipFlag translates a local file path in a pip flag to be absolute.
+	TranslateModeEnvironmentPipFlag
 )

bundle/config/mutator/paths/translation_mode.go

@@ -104,8 +104,10 @@ func (t *translateContext) rewritePath(
 	opts translateOptions,
 ) (string, error) {
 	// If the input is a local requirements file, we need to translate it to an absolute path.
-	if reqPath, ok := libraries.IsLocalRequirementsFile(input); ok {
+	var flag string
+	if reqPath, flagPrefix, ok := libraries.IsLocalPathInPipFlag(input); ok {
 		input = reqPath
+		flag = flagPrefix
 	}
 
 	// We assume absolute paths point to a location in the workspace
@@ -160,10 +162,10 @@ func (t *translateContext) rewritePath(
 		interp, err = t.translateLocalRelativePath(ctx, input, localPath, localRelPath)
 	case paths.TranslateModeLocalRelativeWithPrefix:
 		interp, err = t.translateLocalRelativeWithPrefixPath(ctx, input, localPath, localRelPath)
-	case paths.TranslateModeEnvironmentRequirements:
+	case paths.TranslateModeEnvironmentPipFlag:
 		interp, err = t.translateFilePath(ctx, input, localPath, localRelPath)
-		// Add the -r flag to the path to indicate it's a requirements file used for environment dependencies.
-		interp = "-r " + interp
+		// Add the flag prefix to the path to indicate it's a local file used in pip flags for environment dependencies.
+		interp = flag + " " + interp
 	default:
 		return "", fmt.Errorf("unsupported translate mode: %d", opts.Mode)
 	}

bundle/config/mutator/translate_paths.go

@@ -58,17 +58,11 @@ func IsLibraryLocal(dep string) bool {
 		}
 	}
 
-	// If the dependency starts with --, it's a pip flag option which is a valid
-	// entry for environment dependencies but not a local path
-	if containsPipFlag(dep) {
-		return false
-	}
-
-	// If the dependency is a requirements file, it can either be a local path or a remote path.
-	// Even though the path to requirements.txt can be local we don't return true in this function anyway
+	// If the dependency starts with - or --, it's a pip flag option which is a valid
+	// entry for environment dependencies. Even though the path in the flag can be local we don't return true in this function anyway
 	// and don't treat such path as a local library path.
 	// Instead we handle translation of these paths in a separate code path in TranslatePath mutator.
-	if strings.HasPrefix(dep, "-r") {
+	if containsPipFlag(dep) {
 		return false
 	}
 
@@ -80,18 +74,28 @@ func IsLibraryLocal(dep string) bool {
 	return IsLocalPath(dep)
 }
 
-func IsLocalRequirementsFile(dep string) (string, bool) {
-	dep, ok := strings.CutPrefix(dep, "-r ")
-	if !ok {
-		return "", false
+var PipFlagsWithLocalPaths = []string{
+	"-r",
+	"-e",
+}
+
+func IsLocalPathInPipFlag(dep string) (string, string, bool) {
+	for _, flag := range PipFlagsWithLocalPaths {
+		dep, ok := strings.CutPrefix(dep, flag+" ")
+		if ok {
+			dep = strings.TrimSpace(dep)
+			return dep, flag, IsLocalPath(dep)
+		}
 	}
 
-	dep = strings.TrimSpace(dep)
-	return dep, IsLocalPath(dep)
+	return "", "", false
 }
 
 func containsPipFlag(input string) bool {
-	re := regexp.MustCompile(`--[a-zA-Z0-9-]+`)
+	// Trailing space means the the flag takes an argument or there's multiple arguments in input
+	// Alternatively it could be a flag with no argument and no space after it
+	// For example: -r myfile.txt or --index-url http://myindexurl.com or -i
+	re := regexp.MustCompile(`--?[a-zA-Z0-9-]+(\s|$)`)
 	return re.MatchString(input)
 }