From 662234fb979297edafd19a325600a515272c3189 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Tue, 3 Sep 2024 14:46:52 +0200 Subject: [PATCH 01/13] Fix "bundle init" when run from Databricks --- libs/template/file.go | 63 ++++++++++++++++++++++++++++++++---- libs/template/materialize.go | 7 ++++ 2 files changed, 64 insertions(+), 6 deletions(-) diff --git a/libs/template/file.go b/libs/template/file.go index aafb1acf..d01ab49e 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -2,12 +2,16 @@ package template import ( "context" + "encoding/base64" "io" "io/fs" "os" "path/filepath" + "strings" "github.com/databricks/cli/libs/filer" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/workspace" ) // Interface representing a file to be materialized from a template into a project @@ -53,6 +57,10 @@ type copyFile struct { srcPath string } +func (f *copyFile) isNotebook() bool { + return strings.HasSuffix(f.DstPath().relPath, ".ipynb") +} + func (f *copyFile) DstPath() *destinationPath { return f.dstPath } @@ -68,13 +76,22 @@ 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) - if err != nil { + + if runsOnDatabricks() && f.isNotebook() { + content, err := io.ReadAll(srcFile) + if err != nil { + return err + } + return writeNotebook(path, content) + } else { + dstFile, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, f.perm) + if err != nil { + return err + } + defer dstFile.Close() + _, err = io.Copy(dstFile, srcFile) return err } - defer dstFile.Close() - _, err = io.Copy(dstFile, srcFile) - return err } type inMemoryFile struct { @@ -86,6 +103,10 @@ type inMemoryFile struct { perm fs.FileMode } +func (f *inMemoryFile) isNotebook() bool { + return strings.HasSuffix(f.DstPath().relPath, ".ipynb") +} + func (f *inMemoryFile) DstPath() *destinationPath { return f.dstPath } @@ -97,5 +118,35 @@ func (f *inMemoryFile) PersistToDisk() error { if err != nil { return err } - return os.WriteFile(path, f.content, f.perm) + + if runsOnDatabricks() && f.isNotebook() { + return writeNotebook(path, f.content) + } else { + return os.WriteFile(path, f.content, f.perm) + } +} + +func runsOnDatabricks() bool { + return os.Getenv("DATABRICKS_RUNTIME_VERSION") != "" +} + +func writeNotebook(path string, content []byte) error { + if !strings.HasPrefix(path, "/Workspace/") { + return os.WriteFile(path, content, 0644) + } else { + w, err := databricks.NewWorkspaceClient() + if err != nil { + return err + } + + ctx := context.Background() + + err = w.Workspace.Import(ctx, workspace.Import{ + Format: "AUTO", + Overwrite: false, + Path: path, + Content: base64.StdEncoding.EncodeToString(content), + }) + return err + } } diff --git a/libs/template/materialize.go b/libs/template/materialize.go index d824bf38..bc7da770 100644 --- a/libs/template/materialize.go +++ b/libs/template/materialize.go @@ -40,6 +40,13 @@ func Materialize(ctx context.Context, configFilePath, templateRoot, outputDir st return err } + if outputDir == "" { + outputDir, err = os.Getwd() + if err != nil { + return err + } + } + templatePath := filepath.Join(templateRoot, templateDirName) libraryPath := filepath.Join(templateRoot, libraryDirName) schemaPath := filepath.Join(templateRoot, schemaFileName) From 8d78809469ea414110375967321cf3cbe5207400 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Tue, 3 Sep 2024 16:08:28 +0200 Subject: [PATCH 02/13] use env.lookup --- libs/template/file.go | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/libs/template/file.go b/libs/template/file.go index d01ab49e..acee40f3 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -9,6 +9,7 @@ import ( "path/filepath" "strings" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/workspace" @@ -77,12 +78,14 @@ func (f *copyFile) PersistToDisk() error { } defer srcFile.Close() - if runsOnDatabricks() && f.isNotebook() { + ctx := context.Background() + + if runsOnDatabricks(ctx) && f.isNotebook() { content, err := io.ReadAll(srcFile) if err != nil { return err } - return writeNotebook(path, content) + return writeNotebook(ctx, path, content) } else { dstFile, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, f.perm) if err != nil { @@ -119,18 +122,21 @@ func (f *inMemoryFile) PersistToDisk() error { return err } - if runsOnDatabricks() && f.isNotebook() { - return writeNotebook(path, f.content) + ctx := context.Background() + + if runsOnDatabricks(ctx) && f.isNotebook() { + return writeNotebook(ctx, path, f.content) } else { return os.WriteFile(path, f.content, f.perm) } } -func runsOnDatabricks() bool { - return os.Getenv("DATABRICKS_RUNTIME_VERSION") != "" +func runsOnDatabricks(ctx context.Context) bool { + _, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION") + return ok } -func writeNotebook(path string, content []byte) error { +func writeNotebook(ctx context.Context, path string, content []byte) error { if !strings.HasPrefix(path, "/Workspace/") { return os.WriteFile(path, content, 0644) } else { @@ -139,8 +145,6 @@ func writeNotebook(path string, content []byte) error { return err } - ctx := context.Background() - err = w.Workspace.Import(ctx, workspace.Import{ Format: "AUTO", Overwrite: false, From 6585a757bb9d27eab5698035fe3fbe43f1ebf950 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 09:59:01 +0200 Subject: [PATCH 03/13] address some of the PR feedback --- libs/template/file.go | 73 ++++++++++++++------------------------- libs/template/renderer.go | 2 +- 2 files changed, 26 insertions(+), 49 deletions(-) diff --git a/libs/template/file.go b/libs/template/file.go index acee40f3..fb1303ba 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() error + PersistToDisk(ctx context.Context) error } type destinationPath struct { @@ -58,15 +58,11 @@ type copyFile struct { srcPath string } -func (f *copyFile) isNotebook() bool { - return strings.HasSuffix(f.DstPath().relPath, ".ipynb") -} - func (f *copyFile) DstPath() *destinationPath { return f.dstPath } -func (f *copyFile) PersistToDisk() error { +func (f *copyFile) PersistToDisk(ctx context.Context) error { path := f.DstPath().absPath() err := os.MkdirAll(filepath.Dir(path), 0755) if err != nil { @@ -78,23 +74,11 @@ func (f *copyFile) PersistToDisk() error { } defer srcFile.Close() - ctx := context.Background() - - if runsOnDatabricks(ctx) && f.isNotebook() { - content, err := io.ReadAll(srcFile) - if err != nil { - return err - } - return writeNotebook(ctx, path, content) - } else { - dstFile, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, f.perm) - if err != nil { - return err - } - defer dstFile.Close() - _, err = io.Copy(dstFile, srcFile) + content, err := io.ReadAll(srcFile) + if err != nil { return err } + return writeFile(ctx, path, content) } type inMemoryFile struct { @@ -106,15 +90,11 @@ type inMemoryFile struct { perm fs.FileMode } -func (f *inMemoryFile) isNotebook() bool { - return strings.HasSuffix(f.DstPath().relPath, ".ipynb") -} - func (f *inMemoryFile) DstPath() *destinationPath { return f.dstPath } -func (f *inMemoryFile) PersistToDisk() error { +func (f *inMemoryFile) PersistToDisk(ctx context.Context) error { path := f.DstPath().absPath() err := os.MkdirAll(filepath.Dir(path), 0755) @@ -122,13 +102,7 @@ func (f *inMemoryFile) PersistToDisk() error { return err } - ctx := context.Background() - - if runsOnDatabricks(ctx) && f.isNotebook() { - return writeNotebook(ctx, path, f.content) - } else { - return os.WriteFile(path, f.content, f.perm) - } + return writeFile(ctx, path, f.content) } func runsOnDatabricks(ctx context.Context) bool { @@ -136,21 +110,24 @@ func runsOnDatabricks(ctx context.Context) bool { return ok } -func writeNotebook(ctx context.Context, path string, content []byte) error { - if !strings.HasPrefix(path, "/Workspace/") { - return os.WriteFile(path, content, 0644) +func writeFile(ctx context.Context, path string, content []byte) error { + if strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") { + return importNotebook(ctx, path, content) } else { - w, err := databricks.NewWorkspaceClient() - if err != nil { - return err - } - - err = w.Workspace.Import(ctx, workspace.Import{ - Format: "AUTO", - Overwrite: false, - Path: path, - Content: base64.StdEncoding.EncodeToString(content), - }) - return err + return os.WriteFile(path, content, 0644) } } + +func importNotebook(ctx context.Context, path string, content []byte) error { + w, err := databricks.NewWorkspaceClient() + if err != nil { + return err + } + + return w.Workspace.Import(ctx, workspace.Import{ + Format: "AUTO", + Overwrite: false, + Path: path, + Content: base64.StdEncoding.EncodeToString(content), + }) +} diff --git a/libs/template/renderer.go b/libs/template/renderer.go index 827f3013..4abf99eb 100644 --- a/libs/template/renderer.go +++ b/libs/template/renderer.go @@ -321,7 +321,7 @@ func (r *renderer) persistToDisk() error { // Persist files to disk for _, file := range filesToPersist { - err := file.PersistToDisk() + err := file.PersistToDisk(r.ctx) if err != nil { return err } From 62b2451c1072a804cf21e391f0f0ba28dc6adb67 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 10:31:40 +0200 Subject: [PATCH 04/13] fix tests --- libs/template/file.go | 7 ++----- libs/template/file_test.go | 7 ++++--- libs/template/renderer.go | 3 +-- 3 files changed, 7 insertions(+), 10 deletions(-) diff --git a/libs/template/file.go b/libs/template/file.go index fb1303ba..f635a2de 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -9,9 +9,9 @@ import ( "path/filepath" "strings" + "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/filer" - "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/workspace" ) @@ -119,10 +119,7 @@ func writeFile(ctx context.Context, path string, content []byte) error { } func importNotebook(ctx context.Context, path string, content []byte) error { - w, err := databricks.NewWorkspaceClient() - if err != nil { - return err - } + w := root.WorkspaceClient(ctx) return w.Workspace.Import(ctx, workspace.Import{ Format: "AUTO", diff --git a/libs/template/file_test.go b/libs/template/file_test.go index 85938895..5e3ab395 100644 --- a/libs/template/file_test.go +++ b/libs/template/file_test.go @@ -24,7 +24,7 @@ func testInMemoryFile(t *testing.T, perm fs.FileMode) { perm: perm, content: []byte("123"), } - err := f.PersistToDisk() + err := f.PersistToDisk(context.Background()) assert.NoError(t, err) assertFileContent(t, filepath.Join(tmpDir, "a/b/c"), "123") @@ -38,9 +38,10 @@ 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: context.Background(), + ctx: ctx, dstPath: &destinationPath{ root: tmpDir, relPath: "a/b/c", @@ -49,7 +50,7 @@ func testCopyFile(t *testing.T, perm fs.FileMode) { srcPath: "source", srcFiler: templateFiler, } - err = f.PersistToDisk() + err = f.PersistToDisk(ctx) assert.NoError(t, err) assertFileContent(t, filepath.Join(tmpDir, "a/b/c"), "qwerty") diff --git a/libs/template/renderer.go b/libs/template/renderer.go index 4abf99eb..72ef8ee5 100644 --- a/libs/template/renderer.go +++ b/libs/template/renderer.go @@ -313,8 +313,7 @@ func (r *renderer) persistToDisk() error { _, err := os.Stat(path) 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) { + } else if !errors.Is(err, fs.ErrNotExist) { return fmt.Errorf("error while verifying file %s does not already exist: %w", path, err) } } From 08b6b10ff067de22e90d123a2147f5515ce27b64 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 11:16:50 +0200 Subject: [PATCH 05/13] add unit tests --- libs/template/file.go | 22 +++++++++++------ libs/template/file_test.go | 45 +++++++++++++++++++++++++++++++--- libs/template/renderer.go | 3 ++- libs/template/renderer_test.go | 4 +++ 4 files changed, 61 insertions(+), 13 deletions(-) diff --git a/libs/template/file.go b/libs/template/file.go index f635a2de..8553d15f 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 5e3ab395..9fe3c36f 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 72ef8ee5..3739fc6b 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 92133c5f..25850126 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", From 90d6490106820c1c16c0e6e6abbeb3bfd4c63062 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 11:40:15 +0200 Subject: [PATCH 06/13] Centralize code to detect if we are running on Databricks --- bundle/config/mutator/configure_wsfs.go | 6 ++---- libs/runtime/detect.go | 14 ++++++++++++++ libs/template/file.go | 9 ++------- 3 files changed, 18 insertions(+), 11 deletions(-) create mode 100644 libs/runtime/detect.go diff --git a/bundle/config/mutator/configure_wsfs.go b/bundle/config/mutator/configure_wsfs.go index 1d1bec58..296536d1 100644 --- a/bundle/config/mutator/configure_wsfs.go +++ b/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 } diff --git a/libs/runtime/detect.go b/libs/runtime/detect.go new file mode 100644 index 00000000..cfee420e --- /dev/null +++ b/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 { + _, ok := env.Lookup(ctx, envDatabricksRuntimeVersion) + return ok +} diff --git a/libs/template/file.go b/libs/template/file.go index 8553d15f..c27857c9 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -10,8 +10,8 @@ import ( "strings" "github.com/databricks/cli/cmd/root" - "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/runtime" "github.com/databricks/databricks-sdk-go/service/workspace" ) @@ -107,13 +107,8 @@ func (f *inMemoryFile) PersistToDisk() error { return writeFile(f.ctx, path, f.content, f.perm) } -func runsOnDatabricks(ctx context.Context) bool { - _, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION") - return ok -} - func shouldUseImportNotebook(ctx context.Context, path string) bool { - return strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") + return strings.HasPrefix(path, "/Workspace/") && runtime.RunsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") } func writeFile(ctx context.Context, path string, content []byte, perm fs.FileMode) error { From e51037ff4bbb32017299fd27aec24456d5fe1a9b Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 11:45:27 +0200 Subject: [PATCH 07/13] fixes --- libs/template/materialize.go | 7 ------- libs/template/renderer.go | 9 +++++++-- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/libs/template/materialize.go b/libs/template/materialize.go index bc7da770..d824bf38 100644 --- a/libs/template/materialize.go +++ b/libs/template/materialize.go @@ -40,13 +40,6 @@ func Materialize(ctx context.Context, configFilePath, templateRoot, outputDir st return err } - if outputDir == "" { - outputDir, err = os.Getwd() - if err != nil { - return err - } - } - templatePath := filepath.Join(templateRoot, templateDirName) libraryPath := filepath.Join(templateRoot, libraryDirName) schemaPath := filepath.Join(templateRoot, schemaFileName) diff --git a/libs/template/renderer.go b/libs/template/renderer.go index 3739fc6b..fc4a4d89 100644 --- a/libs/template/renderer.go +++ b/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 ©File{ dstPath: &destinationPath{ - root: r.instanceRoot, + root: rootPath, relPath: relPath, }, perm: perm, @@ -196,7 +201,7 @@ func (r *renderer) computeFile(relPathTemplate string) (file, error) { return &inMemoryFile{ ctx: r.ctx, dstPath: &destinationPath{ - root: r.instanceRoot, + root: rootPath, relPath: relPath, }, perm: perm, From 02745de23712055ce8be60711133c3bb7797647e Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 13:16:00 +0200 Subject: [PATCH 08/13] Use `notebook.Detect` --- libs/notebook/detect.go | 43 ++++++++++++++++++++++++++++++++++++ libs/notebook/detect_test.go | 19 ++++++++++++++++ libs/template/file.go | 12 +++++++--- libs/template/file_test.go | 18 ++++++++------- 4 files changed, 81 insertions(+), 11 deletions(-) diff --git a/libs/notebook/detect.go b/libs/notebook/detect.go index 582a8847..480fae21 100644 --- a/libs/notebook/detect.go +++ b/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) +} diff --git a/libs/notebook/detect_test.go b/libs/notebook/detect_test.go index ad89d6dd..a5892fe7 100644 --- a/libs/notebook/detect_test.go +++ b/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) +} diff --git a/libs/template/file.go b/libs/template/file.go index c27857c9..3f1a3d72 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -11,6 +11,7 @@ import ( "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/notebook" "github.com/databricks/cli/libs/runtime" "github.com/databricks/databricks-sdk-go/service/workspace" ) @@ -107,12 +108,17 @@ func (f *inMemoryFile) PersistToDisk() error { return writeFile(f.ctx, path, f.content, f.perm) } -func shouldUseImportNotebook(ctx context.Context, path string) bool { - return strings.HasPrefix(path, "/Workspace/") && runtime.RunsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") +func shouldUseImportNotebook(ctx context.Context, path string, content []byte) bool { + if strings.HasPrefix(path, "/Workspace/") && runtime.RunsOnDatabricks(ctx) { + isNotebook, _, _ := notebook.DetectWithContent(path, content) + return isNotebook + } else { + return false + } } func writeFile(ctx context.Context, path string, content []byte, perm fs.FileMode) error { - if shouldUseImportNotebook(ctx, path) { + if shouldUseImportNotebook(ctx, path, content) { return importNotebook(ctx, path, content) } else { return os.WriteFile(path, content, perm) diff --git a/libs/template/file_test.go b/libs/template/file_test.go index 9fe3c36f..e1b37e55 100644 --- a/libs/template/file_test.go +++ b/libs/template/file_test.go @@ -118,16 +118,18 @@ func TestTemplateCopyFilePersistToDiskForWindows(t *testing.T) { 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")) + 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")) - 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")) + 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) { From 4c9bf79c76f3ce91106df7f9f179dfd4d782182d Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 13:23:26 +0200 Subject: [PATCH 09/13] fix nit --- libs/template/renderer.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/template/renderer.go b/libs/template/renderer.go index fc4a4d89..30b4db0c 100644 --- a/libs/template/renderer.go +++ b/libs/template/renderer.go @@ -319,7 +319,8 @@ func (r *renderer) persistToDisk() error { _, err := os.Stat(path) if err == nil { return fmt.Errorf("failed to initialize template, one or more files already exist: %s", path) - } else if !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) } } From 694413b115b7aa1691a81686fea641c6da607a3a Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 13:56:36 +0200 Subject: [PATCH 10/13] PR feedback --- libs/template/file.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/libs/template/file.go b/libs/template/file.go index 3f1a3d72..b26435d9 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -11,6 +11,7 @@ import ( "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" @@ -110,11 +111,14 @@ func (f *inMemoryFile) PersistToDisk() error { func shouldUseImportNotebook(ctx context.Context, path string, content []byte) bool { if strings.HasPrefix(path, "/Workspace/") && runtime.RunsOnDatabricks(ctx) { - isNotebook, _, _ := notebook.DetectWithContent(path, content) - return isNotebook - } else { - return false + 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 { From 4cd9b17625c4e63839598d271aca54293dfa1383 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 14:08:11 +0200 Subject: [PATCH 11/13] add test --- libs/runtime/detect.go | 4 ++-- libs/runtime/detect_test.go | 18 ++++++++++++++++++ libs/template/file.go | 2 +- 3 files changed, 21 insertions(+), 3 deletions(-) create mode 100644 libs/runtime/detect_test.go diff --git a/libs/runtime/detect.go b/libs/runtime/detect.go index cfee420e..765eea8d 100644 --- a/libs/runtime/detect.go +++ b/libs/runtime/detect.go @@ -9,6 +9,6 @@ import ( const envDatabricksRuntimeVersion = "DATABRICKS_RUNTIME_VERSION" func RunsOnDatabricks(ctx context.Context) bool { - _, ok := env.Lookup(ctx, envDatabricksRuntimeVersion) - return ok + value, ok := env.Lookup(ctx, envDatabricksRuntimeVersion) + return value != "" && ok } diff --git a/libs/runtime/detect_test.go b/libs/runtime/detect_test.go new file mode 100644 index 00000000..ccc2254e --- /dev/null +++ b/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)) +} diff --git a/libs/template/file.go b/libs/template/file.go index b26435d9..7293ed0a 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -115,7 +115,7 @@ func shouldUseImportNotebook(ctx context.Context, path string, content []byte) b if err != nil { log.Debugf(ctx, "Error detecting notebook: %v", err) } - return isNotebook && err != nil + return isNotebook && err == nil } return false From 84d1bbf2d511d84761865ad929e5a47cd9b62d50 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Wed, 4 Sep 2024 15:14:44 +0200 Subject: [PATCH 12/13] add comment --- libs/template/file.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/libs/template/file.go b/libs/template/file.go index 7293ed0a..826ded37 100644 --- a/libs/template/file.go +++ b/libs/template/file.go @@ -76,6 +76,9 @@ func (f *copyFile) PersistToDisk() error { } defer srcFile.Close() + // 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 From 6bf59ff981343693867d011b574dc5a928fa26fc Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Tue, 17 Sep 2024 11:21:15 +0200 Subject: [PATCH 13/13] address PR feedback --- libs/notebook/detect.go | 15 +++------------ libs/template/renderer.go | 1 + 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/libs/notebook/detect.go b/libs/notebook/detect.go index 480fae21..0d49b3da 100644 --- a/libs/notebook/detect.go +++ b/libs/notebook/detect.go @@ -141,8 +141,7 @@ func Detect(name string) (notebook bool, language workspace.Language, err error) } type inMemoryFile struct { - content []byte - readIndex int64 + buffer bytes.Buffer } type inMemoryFS struct { @@ -158,20 +157,12 @@ func (f *inMemoryFile) Stat() (fs.FileInfo, error) { } 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 + return f.buffer.Read(b) } func (fs inMemoryFS) Open(name string) (fs.File, error) { return &inMemoryFile{ - content: fs.content, - readIndex: 0, + buffer: *bytes.NewBuffer(fs.content), }, nil } diff --git a/libs/template/renderer.go b/libs/template/renderer.go index 30b4db0c..c464df3f 100644 --- a/libs/template/renderer.go +++ b/libs/template/renderer.go @@ -153,6 +153,7 @@ func (r *renderer) computeFile(relPathTemplate string) (file, error) { return nil, err } + // we need the absolute path in case we need to write notebooks using the REST API rootPath, err := filepath.Abs(r.instanceRoot) if err != nil { return nil, err