diff --git a/libs/template/file.go b/libs/template/file.go index f635a2de6b..8553d15f0d 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -23,7 +23,7 @@ type file interface { DstPath() *destinationPath // Write file to disk at the destination path. - PersistToDisk(ctx context.Context) error + PersistToDisk() error } type destinationPath struct { @@ -62,7 +62,7 @@ func (f *copyFile) DstPath() *destinationPath { return f.dstPath } -func (f *copyFile) PersistToDisk(ctx context.Context) error { +func (f *copyFile) PersistToDisk() error { path := f.DstPath().absPath() err := os.MkdirAll(filepath.Dir(path), 0755) if err != nil { @@ -78,10 +78,12 @@ func (f *copyFile) PersistToDisk(ctx context.Context) error { if err != nil { return err } - return writeFile(ctx, path, content) + return writeFile(f.ctx, path, content, f.perm) } type inMemoryFile struct { + ctx context.Context + dstPath *destinationPath content []byte @@ -94,7 +96,7 @@ func (f *inMemoryFile) DstPath() *destinationPath { return f.dstPath } -func (f *inMemoryFile) PersistToDisk(ctx context.Context) error { +func (f *inMemoryFile) PersistToDisk() error { path := f.DstPath().absPath() err := os.MkdirAll(filepath.Dir(path), 0755) @@ -102,7 +104,7 @@ func (f *inMemoryFile) PersistToDisk(ctx context.Context) error { return err } - return writeFile(ctx, path, f.content) + return writeFile(f.ctx, path, f.content, f.perm) } func runsOnDatabricks(ctx context.Context) bool { @@ -110,11 +112,15 @@ func runsOnDatabricks(ctx context.Context) bool { return ok } -func writeFile(ctx context.Context, path string, content []byte) error { - if strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") { +func shouldUseImportNotebook(ctx context.Context, path string) bool { + return strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") +} + +func writeFile(ctx context.Context, path string, content []byte, perm fs.FileMode) error { + if shouldUseImportNotebook(ctx, path) { return importNotebook(ctx, path, content) } else { - return os.WriteFile(path, content, 0644) + return os.WriteFile(path, content, perm) } } diff --git a/libs/template/file_test.go b/libs/template/file_test.go index 5e3ab395d6..9fe3c36fd6 100644 --- a/libs/template/file_test.go +++ b/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", @@ -24,7 +30,7 @@ func testInMemoryFile(t *testing.T, perm fs.FileMode) { perm: perm, content: []byte("123"), } - err := f.PersistToDisk(context.Background()) + err := f.PersistToDisk() assert.NoError(t, err) assertFileContent(t, filepath.Join(tmpDir, "a/b/c"), "123") @@ -38,10 +44,9 @@ func testCopyFile(t *testing.T, perm fs.FileMode) { require.NoError(t, err) err = os.WriteFile(filepath.Join(tmpDir, "source"), []byte("qwerty"), perm) require.NoError(t, err) - ctx := context.Background() f := ©File{ - ctx: ctx, + ctx: context.Background(), dstPath: &destinationPath{ root: tmpDir, relPath: "a/b/c", @@ -50,7 +55,7 @@ func testCopyFile(t *testing.T, perm fs.FileMode) { srcPath: "source", srcFiler: templateFiler, } - err = f.PersistToDisk(ctx) + err = f.PersistToDisk() assert.NoError(t, err) assertFileContent(t, filepath.Join(tmpDir, "a/b/c"), "qwerty") @@ -110,3 +115,35 @@ 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() + assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar")) + assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar.ipynb")) + assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar")) + assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar.ipynb")) + + t.Setenv("DATABRICKS_RUNTIME_VERSION", "14.3") + assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar")) + assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar.ipynb")) + assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar")) + assert.True(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar.ipynb")) +} + +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) +} diff --git a/libs/template/renderer.go b/libs/template/renderer.go index 72ef8ee5a9..3739fc6b6f 100644 --- a/libs/template/renderer.go +++ b/libs/template/renderer.go @@ -194,6 +194,7 @@ func (r *renderer) computeFile(relPathTemplate string) (file, error) { } return &inMemoryFile{ + ctx: r.ctx, dstPath: &destinationPath{ root: r.instanceRoot, relPath: relPath, @@ -320,7 +321,7 @@ func (r *renderer) persistToDisk() error { // Persist files to disk for _, file := range filesToPersist { - err := file.PersistToDisk(r.ctx) + err := file.PersistToDisk() if err != nil { return err } diff --git a/libs/template/renderer_test.go b/libs/template/renderer_test.go index 92133c5fea..25850126a6 100644 --- a/libs/template/renderer_test.go +++ b/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",