Fix "bundle init" when run from Databricks

bundle/config/mutator/configure_wsfs.go

@@ -6,13 +6,11 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/libs/diag"
-	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/cli/libs/runtime"
 	"github.com/databricks/cli/libs/vfs"
 )
 
-const envDatabricksRuntimeVersion = "DATABRICKS_RUNTIME_VERSION"
-
 type configureWSFS struct{}
 
 func ConfigureWSFS() bundle.Mutator {
@@ -32,7 +30,7 @@ func (m *configureWSFS) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno
 	}
 
 	// The executable must be running on DBR.
-	if _, ok := env.Lookup(ctx, envDatabricksRuntimeVersion); !ok {
+	if !runtime.RunsOnDatabricks(ctx) {
 		return nil
 	}
 

libs/notebook/detect.go

@@ -139,3 +139,46 @@ func Detect(name string) (notebook bool, language workspace.Language, err error)
 	b := filepath.Base(name)
 	return DetectWithFS(os.DirFS(d), b)
 }
+
+type inMemoryFile struct {
+	content   []byte
+	readIndex int64
+}
+
+type inMemoryFS struct {
+	content []byte
+}
+
+func (f *inMemoryFile) Close() error {
+	return nil
+}
+
+func (f *inMemoryFile) Stat() (fs.FileInfo, error) {
+	return nil, nil
+}
+
+func (f *inMemoryFile) Read(b []byte) (n int, err error) {
+	if f.readIndex >= int64(len(f.content)) {
+		err = io.EOF
+		return
+	}
+
+	n = copy(b, f.content[f.readIndex:])
+	f.readIndex += int64(n)
+	return
+}
+
+func (fs inMemoryFS) Open(name string) (fs.File, error) {
+	return &inMemoryFile{
+		content:   fs.content,
+		readIndex: 0,
+	}, nil
+}
+
+func DetectWithContent(name string, content []byte) (notebook bool, language workspace.Language, err error) {
+	fs := inMemoryFS{
+		content: content,
+	}
+
+	return DetectWithFS(fs, name)
+}

libs/notebook/detect_test.go

@@ -117,3 +117,22 @@ func TestDetectWithObjectInfo(t *testing.T) {
 	assert.True(t, nb)
 	assert.Equal(t, workspace.LanguagePython, lang)
 }
+
+func TestInMemoryFiles(t *testing.T) {
+	isNotebook, language, err := DetectWithContent("hello.py", []byte("# Databricks notebook source\n print('hello')"))
+	assert.True(t, isNotebook)
+	assert.Equal(t, workspace.LanguagePython, language)
+	require.NoError(t, err)
+
+	isNotebook, language, err = DetectWithContent("hello.py", []byte("print('hello')"))
+	assert.False(t, isNotebook)
+	assert.Equal(t, workspace.Language(""), language)
+	require.NoError(t, err)
+
+	fileContent, err := os.ReadFile("./testdata/py_ipynb.ipynb")
+	require.NoError(t, err)
+	isNotebook, language, err = DetectWithContent("py_ipynb.ipynb", fileContent)
+	assert.True(t, isNotebook)
+	assert.Equal(t, workspace.LanguagePython, language)
+	require.NoError(t, err)
+}

libs/runtime/detect.go

@@ -0,0 +1,14 @@
+package runtime
+
+import (
+	"context"
+
+	"github.com/databricks/cli/libs/env"
+)
+
+const envDatabricksRuntimeVersion = "DATABRICKS_RUNTIME_VERSION"
+
+func RunsOnDatabricks(ctx context.Context) bool {
+	value, ok := env.Lookup(ctx, envDatabricksRuntimeVersion)
+	return value != "" && ok
+}

libs/runtime/detect_test.go

@@ -0,0 +1,18 @@
+package runtime
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRunsOnDatabricks(t *testing.T) {
+	ctx := context.Background()
+
+	t.Setenv("DATABRICKS_RUNTIME_VERSION", "")
+	assert.False(t, RunsOnDatabricks(ctx))
+
+	t.Setenv("DATABRICKS_RUNTIME_VERSION", "14.3")
+	assert.True(t, RunsOnDatabricks(ctx))
+}

libs/template/file.go

@@ -2,12 +2,19 @@ package template
 
 import (
 	"context"
+	"encoding/base64"
 	"io"
 	"io/fs"
 	"os"
 	"path/filepath"
+	"strings"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/notebook"
+	"github.com/databricks/cli/libs/runtime"
+	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
 // Interface representing a file to be materialized from a template into a project
@@ -68,16 +75,20 @@ func (f *copyFile) PersistToDisk() error {
 		return err
 	}
 	defer srcFile.Close()
-	dstFile, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, f.perm)
+
+	// we read the full file into memory because we need to inspect the content
+	// in order to determine if it is a notebook
+	// Once we stop using the workspace API, we can remove this and write in a streaming fashion
+	content, err := io.ReadAll(srcFile)
 	if err != nil {
 		return err
 	}
-	defer dstFile.Close()
-	_, err = io.Copy(dstFile, srcFile)
-	return err
+	return writeFile(f.ctx, path, content, f.perm)
 }
 
 type inMemoryFile struct {
+	ctx context.Context
+
 	dstPath *destinationPath
 
 	content []byte
@@ -97,5 +108,37 @@ func (f *inMemoryFile) PersistToDisk() error {
 	if err != nil {
 		return err
 	}
-	return os.WriteFile(path, f.content, f.perm)
+
+	return writeFile(f.ctx, path, f.content, f.perm)
+}
+
+func shouldUseImportNotebook(ctx context.Context, path string, content []byte) bool {
+	if strings.HasPrefix(path, "/Workspace/") && runtime.RunsOnDatabricks(ctx) {
+		isNotebook, _, err := notebook.DetectWithContent(path, content)
+		if err != nil {
+			log.Debugf(ctx, "Error detecting notebook: %v", err)
+		}
+		return isNotebook && err == nil
+	}
+
+	return false
+}
+
+func writeFile(ctx context.Context, path string, content []byte, perm fs.FileMode) error {
+	if shouldUseImportNotebook(ctx, path, content) {
+		return importNotebook(ctx, path, content)
+	} else {
+		return os.WriteFile(path, content, perm)
+	}
+}
+
+func importNotebook(ctx context.Context, path string, content []byte) error {
+	w := root.WorkspaceClient(ctx)
+
+	return w.Workspace.Import(ctx, workspace.Import{
+		Format:    "AUTO",
+		Overwrite: false,
+		Path:      path,
+		Content:   base64.StdEncoding.EncodeToString(content),
+	})
 }

libs/template/file_test.go

@@ -8,6 +8,11 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/workspace"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -17,6 +22,7 @@ func testInMemoryFile(t *testing.T, perm fs.FileMode) {
 	tmpDir := t.TempDir()
 
 	f := &inMemoryFile{
+		ctx: context.Background(),
 		dstPath: &destinationPath{
 			root:    tmpDir,
 			relPath: "a/b/c",
@@ -109,3 +115,37 @@ func TestTemplateCopyFilePersistToDiskForWindows(t *testing.T) {
 	// fs.FileMode values we can use for different operating systems.
 	testCopyFile(t, 0666)
 }
+
+func TestShouldUseImportNotebook(t *testing.T) {
+	ctx := context.Background()
+	data := []byte("# Databricks notebook source\n print('hello')")
+
+	assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar", data))
+	assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar.ipynb", data))
+	assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar", data))
+	assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar.ipynb", data))
+
+	t.Setenv("DATABRICKS_RUNTIME_VERSION", "14.3")
+	assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar", data))
+	assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar.ipynb", data))
+	assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar", data))
+	assert.True(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar.py", data))
+}
+
+func TestImportNotebook(t *testing.T) {
+	ctx := context.Background()
+
+	m := mocks.NewMockWorkspaceClient(t)
+	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
+
+	workspaceApi := m.GetMockWorkspaceAPI()
+	workspaceApi.EXPECT().Import(mock.Anything, workspace.Import{
+		Content:   "cXdlcnR5", // base64 of "qwerty"
+		Format:    "AUTO",
+		Overwrite: false,
+		Path:      "/Workspace/foo/bar.ipynb",
+	}).Return(nil)
+
+	err := importNotebook(ctx, "/Workspace/foo/bar.ipynb", []byte("qwerty"))
+	assert.NoError(t, err)
+}

libs/template/renderer.go

@@ -153,12 +153,17 @@ func (r *renderer) computeFile(relPathTemplate string) (file, error) {
 		return nil, err
 	}
 
+	rootPath, err := filepath.Abs(r.instanceRoot)
+	if err != nil {
+		return nil, err
+	}
+
 	// If file name does not specify the `.tmpl` extension, then it is copied
 	// over as is, without treating it as a template
 	if !strings.HasSuffix(relPathTemplate, templateExtension) {
 		return &copyFile{
 			dstPath: &destinationPath{
-				root:    r.instanceRoot,
+				root:    rootPath,
 				relPath: relPath,
 			},
 			perm:     perm,
@@ -194,8 +199,9 @@ func (r *renderer) computeFile(relPathTemplate string) (file, error) {
 	}
 
 	return &inMemoryFile{
+		ctx: r.ctx,
 		dstPath: &destinationPath{
-			root:    r.instanceRoot,
+			root:    rootPath,
 			relPath: relPath,
 		},
 		perm:    perm,
@@ -314,7 +320,7 @@ func (r *renderer) persistToDisk() error {
 		if err == nil {
 			return fmt.Errorf("failed to initialize template, one or more files already exist: %s", path)
 		}
-		if err != nil && !errors.Is(err, fs.ErrNotExist) {
+		if !errors.Is(err, fs.ErrNotExist) {
 			return fmt.Errorf("error while verifying file %s does not already exist: %w", path, err)
 		}
 	}

libs/template/renderer_test.go

@@ -329,6 +329,7 @@ func TestRendererPersistToDisk(t *testing.T) {
 		skipPatterns: []string{"a/b/c", "mn*"},
 		files: []file{
 			&inMemoryFile{
+				ctx: ctx,
 				dstPath: &destinationPath{
 					root:    tmpDir,
 					relPath: "a/b/c",
@@ -337,6 +338,7 @@ func TestRendererPersistToDisk(t *testing.T) {
 				content: nil,
 			},
 			&inMemoryFile{
+				ctx: ctx,
 				dstPath: &destinationPath{
 					root:    tmpDir,
 					relPath: "mno",
@@ -345,6 +347,7 @@ func TestRendererPersistToDisk(t *testing.T) {
 				content: nil,
 			},
 			&inMemoryFile{
+				ctx: ctx,
 				dstPath: &destinationPath{
 					root:    tmpDir,
 					relPath: "a/b/d",
@@ -353,6 +356,7 @@ func TestRendererPersistToDisk(t *testing.T) {
 				content: []byte("123"),
 			},
 			&inMemoryFile{
+				ctx: ctx,
 				dstPath: &destinationPath{
 					root:    tmpDir,
 					relPath: "mmnn",