steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

address some of the PR feedback

This commit is contained in:
Fabian Jakobs 2024-09-04 09:59:01 +02:00
parent 8d78809469
commit 6585a757bb
No known key found for this signature in database
2 changed files with 26 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
libs/template/file.go
View File

@ -23,7 +23,7 @@ type file interface {
DstPath() *destinationPath DstPath() *destinationPath
// Write file to disk at the destination path. // Write file to disk at the destination path.
PersistToDisk() error PersistToDisk(ctx context.Context) error
} }
type destinationPath struct { type destinationPath struct {
@ -58,15 +58,11 @@ type copyFile struct {
srcPath string srcPath string
} }
func (f *copyFile) isNotebook() bool {
return strings.HasSuffix(f.DstPath().relPath, ".ipynb")
}
func (f *copyFile) DstPath() *destinationPath { func (f *copyFile) DstPath() *destinationPath {
return f.dstPath return f.dstPath
} }
func (f *copyFile) PersistToDisk() error { func (f *copyFile) PersistToDisk(ctx context.Context) error {
path := f.DstPath().absPath() path := f.DstPath().absPath()
err := os.MkdirAll(filepath.Dir(path), 0755) err := os.MkdirAll(filepath.Dir(path), 0755)
if err != nil { if err != nil {
@ -78,23 +74,11 @@ func (f *copyFile) PersistToDisk() error {
} }
defer srcFile.Close() defer srcFile.Close()
ctx := context.Background()
if runsOnDatabricks(ctx) && f.isNotebook() {
content, err := io.ReadAll(srcFile) content, err := io.ReadAll(srcFile)
if err != nil { if err != nil {
return err return err
} }
return writeNotebook(ctx, path, content) return writeFile(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)
return err
}
} }
type inMemoryFile struct { type inMemoryFile struct {
@ -106,15 +90,11 @@ type inMemoryFile struct {
perm fs.FileMode perm fs.FileMode
} }
func (f *inMemoryFile) isNotebook() bool {
return strings.HasSuffix(f.DstPath().relPath, ".ipynb")
}
func (f *inMemoryFile) DstPath() *destinationPath { func (f *inMemoryFile) DstPath() *destinationPath {
return f.dstPath return f.dstPath
} }
func (f *inMemoryFile) PersistToDisk() error { func (f *inMemoryFile) PersistToDisk(ctx context.Context) error {
path := f.DstPath().absPath() path := f.DstPath().absPath()
err := os.MkdirAll(filepath.Dir(path), 0755) err := os.MkdirAll(filepath.Dir(path), 0755)
@ -122,13 +102,7 @@ func (f *inMemoryFile) PersistToDisk() error {
return err return err
} }
ctx := context.Background() return writeFile(ctx, path, f.content)
if runsOnDatabricks(ctx) && f.isNotebook() {
return writeNotebook(ctx, path, f.content)
} else {
return os.WriteFile(path, f.content, f.perm)
}
} }
func runsOnDatabricks(ctx context.Context) bool { func runsOnDatabricks(ctx context.Context) bool {
@ -136,21 +110,24 @@ func runsOnDatabricks(ctx context.Context) bool {
return ok return ok
} }
func writeNotebook(ctx context.Context, path string, content []byte) error { func writeFile(ctx context.Context, path string, content []byte) error {
if !strings.HasPrefix(path, "/Workspace/") { if strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") {
return os.WriteFile(path, content, 0644) return importNotebook(ctx, path, content)
} else { } else {
return os.WriteFile(path, content, 0644)
}
}
func importNotebook(ctx context.Context, path string, content []byte) error {
w, err := databricks.NewWorkspaceClient() w, err := databricks.NewWorkspaceClient()
if err != nil { if err != nil {
return err return err
} }
err = w.Workspace.Import(ctx, workspace.Import{ return w.Workspace.Import(ctx, workspace.Import{
Format: "AUTO", Format: "AUTO",
Overwrite: false, Overwrite: false,
Path: path, Path: path,
Content: base64.StdEncoding.EncodeToString(content), Content: base64.StdEncoding.EncodeToString(content),
}) })
return err
}
} }

2
libs/template/renderer.go
View File

@ -321,7 +321,7 @@ func (r *renderer) persistToDisk() error {
// Persist files to disk // Persist files to disk
for _, file := range filesToPersist { for _, file := range filesToPersist {
err := file.PersistToDisk() err := file.PersistToDisk(r.ctx)
if err != nil { if err != nil {
return err return err
} }