From d0385a0205cd09eb4221f9aa56bf82df7839d955 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Fri, 29 Nov 2024 02:45:33 +0100
Subject: [PATCH] add extractVolumeFromPath

---
 bundle/libraries/filer_volume.go      | 52 ++++++++++++++++++---------
 bundle/libraries/filer_volume_test.go | 30 +++++++++++++---
 2 files changed, 62 insertions(+), 20 deletions(-)

diff --git a/bundle/libraries/filer_volume.go b/bundle/libraries/filer_volume.go
index 7128f4d7..e6b6c0fe 100644
--- a/bundle/libraries/filer_volume.go
+++ b/bundle/libraries/filer_volume.go
@@ -13,6 +13,31 @@ import (
 	"github.com/databricks/cli/libs/filer"
 )
 
+func extractVolumeFromPath(artifactPath string) (string, string, string, error) {
+	if !IsVolumesPath(artifactPath) {
+		return "", "", "", fmt.Errorf("expected artifact_path to start with /Volumes/, got %s", artifactPath)
+	}
+
+	parts := strings.Split(artifactPath, "/")
+	volumeFormatErr := fmt.Errorf("expected UC volume path to be in the format /Volumes/<catalog>/<schema>/<volume>/..., got %s", artifactPath)
+
+	// Incorrect format.
+	if len(parts) < 5 {
+		return "", "", "", volumeFormatErr
+	}
+
+	catalogName := parts[2]
+	schemaName := parts[3]
+	volumeName := parts[4]
+
+	// Incorrect format.
+	if catalogName == "" || schemaName == "" || volumeName == "" {
+		return "", "", "", volumeFormatErr
+	}
+
+	return catalogName, schemaName, volumeName, nil
+}
+
 // This function returns a filer for ".internal" folder inside the directory configured
 // at `workspace.artifact_path`.
 // This function also checks if the UC volume exists in the workspace and then:
@@ -32,26 +57,21 @@ func filerForVolume(ctx context.Context, b *bundle.Bundle) (filer.Filer, string,
 		return nil, "", diag.Errorf("expected artifact_path to start with /Volumes/, got %s", artifactPath)
 	}
 
-	parts := strings.Split(artifactPath, "/")
-	volumeFormatErr := fmt.Errorf("expected UC volume path to be in the format /Volumes/<catalog>/<schema>/<volume>/..., got %s", artifactPath)
-
-	// Incorrect format.
-	if len(parts) < 5 {
-		return nil, "", diag.FromErr(volumeFormatErr)
-	}
-
-	catalogName := parts[2]
-	schemaName := parts[3]
-	volumeName := parts[4]
-
-	// Incorrect format.
-	if catalogName == "" || schemaName == "" || volumeName == "" {
-		return nil, "", diag.FromErr(volumeFormatErr)
+	catalogName, schemaName, volumeName, err := extractVolumeFromPath(artifactPath)
+	if err != nil {
+		return nil, "", diag.Diagnostics{
+			{
+				Severity:  diag.Error,
+				Summary:   err.Error(),
+				Locations: b.Config.GetLocations("workspace.artifact_path"),
+				Paths:     []dyn.Path{dyn.MustPathFromString("workspace.artifact_path")},
+			},
+		}
 	}
 
 	// Check if the UC volume exists in the workspace.
 	volumePath := fmt.Sprintf("/Volumes/%s/%s/%s", catalogName, schemaName, volumeName)
-	err := w.Files.GetDirectoryMetadataByDirectoryPath(ctx, volumePath)
+	err = w.Files.GetDirectoryMetadataByDirectoryPath(ctx, volumePath)
 
 	// If the volume exists already, directly return the filer for the path to
 	// upload the artifacts to.
diff --git a/bundle/libraries/filer_volume_test.go b/bundle/libraries/filer_volume_test.go
index 987c57cb..53f8665a 100644
--- a/bundle/libraries/filer_volume_test.go
+++ b/bundle/libraries/filer_volume_test.go
@@ -153,8 +153,8 @@ func TestFilerForVolumeInBundle(t *testing.T) {
 	})
 }
 
-func TestFilerForVolumeWithInvalidVolumePaths(t *testing.T) {
-	invalidPaths := []string{
+func invalidVolumePaths() []string {
+	return []string{
 		"/Volumes/",
 		"/Volumes/main",
 		"/Volumes/main/",
@@ -165,8 +165,10 @@ func TestFilerForVolumeWithInvalidVolumePaths(t *testing.T) {
 		"/Volumes/main/my_schema//",
 		"/Volumes//my_schema/my_volume",
 	}
+}
 
-	for _, p := range invalidPaths {
+func TestFilerForVolumeWithInvalidVolumePaths(t *testing.T) {
+	for _, p := range invalidVolumePaths() {
 		b := &bundle.Bundle{
 			Config: config.Root{
 				Workspace: config.Workspace{
@@ -175,8 +177,15 @@ func TestFilerForVolumeWithInvalidVolumePaths(t *testing.T) {
 			},
 		}
 
+		bundletest.SetLocation(b, "workspace.artifact_path", []dyn.Location{{File: "config.yml", Line: 1, Column: 2}})
+
 		_, _, diags := GetFilerForLibraries(context.Background(), b)
-		require.EqualError(t, diags.Error(), fmt.Sprintf("expected UC volume path to be in the format /Volumes/<catalog>/<schema>/<volume>/..., got %s", p))
+		require.Equal(t, diags, diag.Diagnostics{{
+			Severity:  diag.Error,
+			Summary:   fmt.Sprintf("expected UC volume path to be in the format /Volumes/<catalog>/<schema>/<volume>/..., got %s", p),
+			Locations: []dyn.Location{{File: "config.yml", Line: 1, Column: 2}},
+			Paths:     []dyn.Path{dyn.MustPathFromString("workspace.artifact_path")},
+		}})
 	}
 }
 
@@ -221,3 +230,16 @@ func TestFilerForVolumeWithValidlVolumePaths(t *testing.T) {
 		assert.IsType(t, &filer.FilesClient{}, client)
 	}
 }
+
+func TestExtractVolumeFromPath(t *testing.T) {
+	catalogName, schemaName, volumeName, err := extractVolumeFromPath("/Volumes/main/my_schema/my_volume")
+	require.NoError(t, err)
+	assert.Equal(t, "main", catalogName)
+	assert.Equal(t, "my_schema", schemaName)
+	assert.Equal(t, "my_volume", volumeName)
+
+	for _, p := range invalidVolumePaths() {
+		_, _, _, err := extractVolumeFromPath(p)
+		assert.EqualError(t, err, fmt.Sprintf("expected UC volume path to be in the format /Volumes/<catalog>/<schema>/<volume>/..., got %s", p))
+	}
+}