Merge branch 'main' of github.com:databricks/cli into tmp/in-place-prototype

2024-11-05 18:13:27 +01:00 · 2024-11-05 18:13:27 +01:00 · ae26beabfe
parent 6b22fdebba b81008e2f6
commit ae26beabfe
17 changed files with 283 additions and 19 deletions
--- a/bundle/config/mutator/translate_paths.go
+++ b/bundle/config/mutator/translate_paths.go
@ -177,7 +177,7 @@ func (t *translateContext) translateNoOp(literal, localFullPath, localRelPath, r
 }

 func (t *translateContext) retainLocalAbsoluteFilePath(literal, localFullPath, localRelPath, remotePath string) (string, error) {
-	info, err := t.b.SyncRoot.Stat(localRelPath)
+	info, err := t.b.SyncRoot.Stat(filepath.ToSlash(localRelPath))
 	if errors.Is(err, fs.ErrNotExist) {
 		return "", fmt.Errorf("file %s not found", literal)
 	}
--- a/bundle/config/mutator/translate_paths_dashboards_test.go
+++ b/bundle/config/mutator/translate_paths_dashboards_test.go
@ -0,0 +1,54 @@
+package mutator_test
+
+import (
+	"context"
+	"path/filepath"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/config/mutator"
+	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/cli/bundle/internal/bundletest"
+	"github.com/databricks/cli/libs/dyn"
+	"github.com/databricks/cli/libs/vfs"
+	"github.com/databricks/databricks-sdk-go/service/dashboards"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTranslatePathsDashboards_FilePathRelativeSubDirectory(t *testing.T) {
+	dir := t.TempDir()
+	touchEmptyFile(t, filepath.Join(dir, "src", "my_dashboard.lvdash.json"))
+
+	b := &bundle.Bundle{
+		SyncRootPath: dir,
+		SyncRoot:     vfs.MustNew(dir),
+		Config: config.Root{
+			Resources: config.Resources{
+				Dashboards: map[string]*resources.Dashboard{
+					"dashboard": {
+						CreateDashboardRequest: &dashboards.CreateDashboardRequest{
+							DisplayName: "My Dashboard",
+						},
+						FilePath: "../src/my_dashboard.lvdash.json",
+					},
+				},
+			},
+		},
+	}
+
+	bundletest.SetLocation(b, "resources.dashboards", []dyn.Location{{
+		File: filepath.Join(dir, "resources/dashboard.yml"),
+	}})
+
+	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
+	require.NoError(t, diags.Error())
+
+	// Assert that the file path for the dashboard has been converted to its local absolute path.
+	assert.Equal(
+		t,
+		filepath.Join(dir, "src", "my_dashboard.lvdash.json"),
+		b.Config.Resources.Dashboards["dashboard"].FilePath,
+	)
+}
--- a/cmd/bundle/run.go
+++ b/cmd/bundle/run.go
@ -35,17 +35,23 @@ func promptRunArgument(ctx context.Context, b *bundle.Bundle) (string, error) {
 	return key, nil
 }

-func resolveRunArgument(ctx context.Context, b *bundle.Bundle, args []string) (string, error) {
+// resolveRunArgument resolves the resource key to run.
+// It returns the remaining arguments to pass to the runner, if applicable.
+func resolveRunArgument(ctx context.Context, b *bundle.Bundle, args []string) (string, []string, error) {
 	// If no arguments are specified, prompt the user to select something to run.
 	if len(args) == 0 && cmdio.IsPromptSupported(ctx) {
-		return promptRunArgument(ctx, b)
+		key, err := promptRunArgument(ctx, b)
+		if err != nil {
+			return "", nil, err
+		}
+		return key, args, nil
 	}

 	if len(args) < 1 {
-		return "", fmt.Errorf("expected a KEY of the resource to run")
+		return "", nil, fmt.Errorf("expected a KEY of the resource to run")
 	}

-	return args[0], nil
+	return args[0], args[1:], nil
 }

 func keyToRunner(b *bundle.Bundle, arg string) (run.Runner, error) {
@ -109,7 +115,7 @@ task or a Python wheel task, the second example applies.
 			return err
 		}

-		arg, err := resolveRunArgument(ctx, b, args)
+		key, args, err := resolveRunArgument(ctx, b, args)
 		if err != nil {
 			return err
 		}
@ -124,13 +130,13 @@ task or a Python wheel task, the second example applies.
 			return err
 		}

-		runner, err := keyToRunner(b, arg)
+		runner, err := keyToRunner(b, key)
 		if err != nil {
 			return err
 		}

 		// Parse additional positional arguments.
-		err = runner.ParseArgs(args[1:], &runOptions)
+		err = runner.ParseArgs(args, &runOptions)
 		if err != nil {
 			return err
 		}
--- a/cmd/root/root.go
+++ b/cmd/root/root.go
@ -75,6 +75,7 @@ func New(ctx context.Context) *cobra.Command {

 		// Configure our user agent with the command that's about to be executed.
 		ctx = withCommandInUserAgent(ctx, cmd)
+		ctx = withCommandExecIdInUserAgent(ctx)
 		ctx = withUpstreamInUserAgent(ctx)
 		cmd.SetContext(ctx)
 		return nil
--- a/cmd/root/user_agent_command_exec_id.go
+++ b/cmd/root/user_agent_command_exec_id.go
@ -0,0 +1,14 @@
+package root
+
+import (
+	"context"
+
+	"github.com/databricks/databricks-sdk-go/useragent"
+	"github.com/google/uuid"
+)
+
+func withCommandExecIdInUserAgent(ctx context.Context) context.Context {
+	// A UUID that will allow us to correlate multiple API requests made by
+	// the same CLI invocation.
+	return useragent.InContext(ctx, "cmd-exec-id", uuid.New().String())
+}
--- a/cmd/root/user_agent_command_exec_id_test.go
+++ b/cmd/root/user_agent_command_exec_id_test.go
@ -0,0 +1,26 @@
+package root
+
+import (
+	"context"
+	"regexp"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/useragent"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithCommandExecIdInUserAgent(t *testing.T) {
+	ctx := withCommandExecIdInUserAgent(context.Background())
+
+	// Check that the command exec ID is in the user agent string.
+	ua := useragent.FromContext(ctx)
+	re := regexp.MustCompile(`cmd-exec-id/([a-f0-9-]+)`)
+	matches := re.FindAllStringSubmatch(ua, -1)
+
+	// Assert that we have exactly one match and that it's a valid UUID.
+	require.Len(t, matches, 1)
+	_, err := uuid.Parse(matches[0][1])
+	assert.NoError(t, err)
+}
--- a/cmd/root/user_agent_command_test.go
+++ b/cmd/root/user_agent_command_test.go
@ -1,13 +1,15 @@
 package root

 import (
+	"context"
 	"testing"

+	"github.com/databricks/databricks-sdk-go/useragent"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 )

-func TestCommandString(t *testing.T) {
+func TestWithCommandInUserAgent(t *testing.T) {
 	root := &cobra.Command{
 		Use: "root",
 	}
@ -26,4 +28,9 @@ func TestCommandString(t *testing.T) {
 	assert.Equal(t, "root", commandString(root))
 	assert.Equal(t, "hello", commandString(hello))
 	assert.Equal(t, "hello_world", commandString(world))
+
+	ctx := withCommandInUserAgent(context.Background(), world)
+
+	ua := useragent.FromContext(ctx)
+	assert.Contains(t, ua, "cmd/hello_world")
 }
--- a/go.mod
+++ b/go.mod
@ -14,7 +14,7 @@ require (
 	github.com/hashicorp/go-version v1.7.0 // MPL 2.0
 	github.com/hashicorp/hc-install v0.9.0 // MPL 2.0
 	github.com/hashicorp/terraform-exec v0.21.0 // MPL 2.0
-	github.com/hashicorp/terraform-json v0.22.1 // MPL 2.0
+	github.com/hashicorp/terraform-json v0.23.0 // MPL 2.0
 	github.com/manifoldco/promptui v0.9.0 // BSD-3-Clause
 	github.com/mattn/go-isatty v0.0.20 // MIT
 	github.com/nwidger/jsoncolor v0.3.2 // MIT
@ -56,7 +56,7 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
-	github.com/zclconf/go-cty v1.14.4 // indirect
+	github.com/zclconf/go-cty v1.15.0 // indirect
 	go.opencensus.io v0.24.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect
 	go.opentelemetry.io/otel v1.24.0 // indirect
--- a/go.sum
+++ b/go.sum
@ -109,8 +109,8 @@ github.com/hashicorp/hc-install v0.9.0 h1:2dIk8LcvANwtv3QZLckxcjyF5w8KVtiMxu6G6e
 github.com/hashicorp/hc-install v0.9.0/go.mod h1:+6vOP+mf3tuGgMApVYtmsnDoKWMDcFXeTxCACYZ8SFg=
 github.com/hashicorp/terraform-exec v0.21.0 h1:uNkLAe95ey5Uux6KJdua6+cv8asgILFVWkd/RG0D2XQ=
 github.com/hashicorp/terraform-exec v0.21.0/go.mod h1:1PPeMYou+KDUSSeRE9szMZ/oHf4fYUmB923Wzbq1ICg=
-github.com/hashicorp/terraform-json v0.22.1 h1:xft84GZR0QzjPVWs4lRUwvTcPnegqlyS7orfb5Ltvec=
-github.com/hashicorp/terraform-json v0.22.1/go.mod h1:JbWSQCLFSXFFhg42T7l9iJwdGXBYV8fmmD6o/ML4p3A=
+github.com/hashicorp/terraform-json v0.23.0 h1:sniCkExU4iKtTADReHzACkk8fnpQXrdD2xoR+lppBkI=
+github.com/hashicorp/terraform-json v0.23.0/go.mod h1:MHdXbBAbSg0GvzuWazEGKAn/cyNfIB7mN6y7KJN6y2c=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOlocH6Fxy8MmwDt+yVQYULKfN0RoTN8A=
@ -160,8 +160,8 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/xanzy/ssh-agent v0.3.3 h1:+/15pJfg/RsTxqYcX6fHqOXZwwMP+2VyYWJeWM2qQFM=
 github.com/xanzy/ssh-agent v0.3.3/go.mod h1:6dzNDKs0J9rVPHPhaGCukekBHKqfl+L3KghI1Bc68Uw=
-github.com/zclconf/go-cty v1.14.4 h1:uXXczd9QDGsgu0i/QFR/hzI5NYCHLf6NQw/atrbnhq8=
-github.com/zclconf/go-cty v1.14.4/go.mod h1:VvMs5i0vgZdhYawQNq5kePSpLAoz8u1xvZgrPIxfnZE=
+github.com/zclconf/go-cty v1.15.0 h1:tTCRWxsexYUmtt/wVxgDClUe+uQusuI443uL6e+5sXQ=
+github.com/zclconf/go-cty v1.15.0/go.mod h1:VvMs5i0vgZdhYawQNq5kePSpLAoz8u1xvZgrPIxfnZE=
 go.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0=
 go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo=
 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 h1:4Pp6oUg3+e/6M4C0A/3kJ2VYa++dsWVTtGgLVj5xtHg=
--- a/internal/bundle/bundles/python_wheel_task_with_cluster/databricks_template_schema.json
+++ b/internal/bundle/bundles/python_wheel_task_with_cluster/databricks_template_schema.json
@ -0,0 +1,25 @@
+{
+    "properties": {
+        "project_name": {
+            "type": "string",
+            "default": "my_test_code",
+            "description": "Unique name for this project"
+        },
+        "spark_version": {
+            "type": "string",
+            "description": "Spark version used for job cluster"
+        },
+        "node_type_id": {
+            "type": "string",
+            "description": "Node type id for job cluster"
+        },
+        "unique_id": {
+            "type": "string",
+            "description": "Unique ID for job name"
+        },
+        "instance_pool_id": {
+            "type": "string",
+            "description": "Instance pool id for job cluster"
+        }
+    }
+}
--- a/internal/bundle/bundles/python_wheel_task_with_cluster/template/databricks.yml.tmpl
+++ b/internal/bundle/bundles/python_wheel_task_with_cluster/template/databricks.yml.tmpl
@ -0,0 +1,29 @@
+bundle:
+  name: wheel-task
+
+workspace:
+  root_path: "~/.bundle/{{.unique_id}}"
+
+resources:
+  clusters:
+    test_cluster:
+      cluster_name: "test-cluster-{{.unique_id}}"
+      spark_version: "{{.spark_version}}"
+      node_type_id: "{{.node_type_id}}"
+      num_workers: 1
+      data_security_mode: USER_ISOLATION
+
+  jobs:
+    some_other_job:
+      name: "[${bundle.target}] Test Wheel Job {{.unique_id}}"
+      tasks:
+        - task_key: TestTask
+          existing_cluster_id: "${resources.clusters.test_cluster.cluster_id}"
+          python_wheel_task:
+            package_name: my_test_code
+            entry_point: run
+            parameters:
+              - "one"
+              - "two"
+          libraries:
+            - whl: ./dist/*.whl
--- a/internal/bundle/bundles/python_wheel_task_with_cluster/template/setup.py.tmpl
+++ b/internal/bundle/bundles/python_wheel_task_with_cluster/template/setup.py.tmpl
@ -0,0 +1,15 @@
+from setuptools import setup, find_packages
+
+import {{.project_name}}
+
+setup(
+    name="{{.project_name}}",
+    version={{.project_name}}.__version__,
+    author={{.project_name}}.__author__,
+    url="https://databricks.com",
+    author_email="john.doe@databricks.com",
+    description="my example wheel",
+    packages=find_packages(include=["{{.project_name}}"]),
+    entry_points={"group1": "run={{.project_name}}.__main__:main"},
+    install_requires=["setuptools"],
+)
--- a/internal/bundle/bundles/python_wheel_task_with_cluster/template/{{.project_name}}/init.py
+++ b/internal/bundle/bundles/python_wheel_task_with_cluster/template/{{.project_name}}/init.py
@ -0,0 +1,2 @@
+__version__ = "0.0.1"
+__author__ = "Databricks"
--- a/internal/bundle/bundles/python_wheel_task_with_cluster/template/{{.project_name}}/main.py
+++ b/internal/bundle/bundles/python_wheel_task_with_cluster/template/{{.project_name}}/main.py
@ -0,0 +1,16 @@
+"""
+The entry point of the Python Wheel
+"""
+
+import sys
+
+
+def main():
+    # This method will print the provided arguments
+    print("Hello from my func")
+    print("Got arguments:")
+    print(sys.argv)
+
+
+if __name__ == "__main__":
+    main()
--- a/internal/bundle/python_wheel_test.go
+++ b/internal/bundle/python_wheel_test.go
@ -5,17 +5,18 @@ import (

 	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/env"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 )

-func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bool) {
+func runPythonWheelTest(t *testing.T, templateName string, sparkVersion string, pythonWheelWrapper bool) {
 	ctx, _ := acc.WorkspaceTest(t)

 	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
 	instancePoolId := env.Get(ctx, "TEST_INSTANCE_POOL_ID")
-	bundleRoot, err := initTestTemplate(t, ctx, "python_wheel_task", map[string]any{
+	bundleRoot, err := initTestTemplate(t, ctx, templateName, map[string]any{
 		"node_type_id":         nodeTypeId,
 		"unique_id":            uuid.New().String(),
 		"spark_version":        sparkVersion,
@ -45,9 +46,19 @@ func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bo
 }

 func TestAccPythonWheelTaskDeployAndRunWithoutWrapper(t *testing.T) {
-	runPythonWheelTest(t, "13.3.x-snapshot-scala2.12", false)
+	runPythonWheelTest(t, "python_wheel_task", "13.3.x-snapshot-scala2.12", false)
 }

 func TestAccPythonWheelTaskDeployAndRunWithWrapper(t *testing.T) {
-	runPythonWheelTest(t, "12.2.x-scala2.12", true)
+	runPythonWheelTest(t, "python_wheel_task", "12.2.x-scala2.12", true)
+}
+
+func TestAccPythonWheelTaskDeployAndRunOnInteractiveCluster(t *testing.T) {
+	_, wt := acc.WorkspaceTest(t)
+
+	if testutil.IsAWSCloud(wt.T) {
+		t.Skip("Skipping test for AWS cloud because it is not permitted to create clusters")
+	}
+
+	runPythonWheelTest(t, "python_wheel_task_with_cluster", defaultSparkVersion, false)
 }
--- a/libs/auth/oauth.go
+++ b/libs/auth/oauth.go
@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/url"
 	"strings"
 	"time"

@ -143,6 +144,26 @@ func (a *PersistentAuth) Challenge(ctx context.Context) error {
 	return nil
 }

+// This function cleans up the host URL by only retaining the scheme and the host.
+// This function thus removes any path, query arguments, or fragments from the URL.
+func (a *PersistentAuth) cleanHost() {
+	parsedHost, err := url.Parse(a.Host)
+	if err != nil {
+		return
+	}
+	// when either host or scheme is empty, we don't want to clean it. This is because
+	// the Go url library parses a raw "abc" string as the path of a URL and cleaning
+	// it will return thus return an empty string.
+	if parsedHost.Host == "" || parsedHost.Scheme == "" {
+		return
+	}
+	host := url.URL{
+		Scheme: parsedHost.Scheme,
+		Host:   parsedHost.Host,
+	}
+	a.Host = host.String()
+}
+
 func (a *PersistentAuth) init(ctx context.Context) error {
 	if a.Host == "" && a.AccountID == "" {
 		return ErrFetchCredentials
@ -156,6 +177,9 @@ func (a *PersistentAuth) init(ctx context.Context) error {
 	if a.browser == nil {
 		a.browser = browser.OpenURL
 	}
+
+	a.cleanHost()
+
 	// try acquire listener, which we also use as a machine-local
 	// exclusive lock to prevent token cache corruption in the scope
 	// of developer machine, where this command runs.
--- a/libs/auth/oauth_test.go
+++ b/libs/auth/oauth_test.go
@ -228,3 +228,37 @@ func TestChallengeFailed(t *testing.T) {
 		assert.EqualError(t, err, "authorize: access_denied: Policy evaluation failed for this request")
 	})
 }
+
+func TestPersistentAuthCleanHost(t *testing.T) {
+	for _, tcases := range []struct {
+		in  string
+		out string
+	}{
+		{"https://example.com", "https://example.com"},
+		{"https://example.com/", "https://example.com"},
+		{"https://example.com/path", "https://example.com"},
+		{"https://example.com/path/subpath", "https://example.com"},
+		{"https://example.com/path?query=1", "https://example.com"},
+		{"https://example.com/path?query=1&other=2", "https://example.com"},
+		{"https://example.com/path#fragment", "https://example.com"},
+		{"https://example.com/path?query=1#fragment", "https://example.com"},
+		{"https://example.com/path?query=1&other=2#fragment", "https://example.com"},
+		{"https://example.com/path/subpath?query=1", "https://example.com"},
+		{"https://example.com/path/subpath?query=1&other=2", "https://example.com"},
+		{"https://example.com/path/subpath#fragment", "https://example.com"},
+		{"https://example.com/path/subpath?query=1#fragment", "https://example.com"},
+		{"https://example.com/path/subpath?query=1&other=2#fragment", "https://example.com"},
+		{"https://example.com/path?query=1%20value&other=2%20value", "https://example.com"},
+		{"http://example.com/path/subpath?query=1%20value&other=2%20value", "http://example.com"},
+
+		// URLs without scheme should be left as is
+		{"abc", "abc"},
+		{"abc.com/def", "abc.com/def"},
+	} {
+		p := &PersistentAuth{
+			Host: tcases.in,
+		}
+		p.cleanHost()
+		assert.Equal(t, tcases.out, p.Host)
+	}
+}