diff --git a/libs/filer/workspace_files_client.go b/libs/filer/workspace_files_client.go index 09f11b161e..d799c1f885 100644 --- a/libs/filer/workspace_files_client.go +++ b/libs/filer/workspace_files_client.go @@ -19,6 +19,7 @@ import ( "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/marshal" "github.com/databricks/databricks-sdk-go/service/workspace" ) @@ -38,7 +39,7 @@ func (entry wsfsDirEntry) Info() (fs.FileInfo, error) { func wsfsDirEntriesFromObjectInfos(objects []workspace.ObjectInfo) []fs.DirEntry { info := make([]fs.DirEntry, len(objects)) for i, v := range objects { - info[i] = wsfsDirEntry{wsfsFileInfo{oi: v}} + info[i] = wsfsDirEntry{wsfsFileInfo{ObjectInfo: v}} } // Sort by name for parity with os.ReadDir. @@ -48,19 +49,22 @@ func wsfsDirEntriesFromObjectInfos(objects []workspace.ObjectInfo) []fs.DirEntry // Type that implements fs.FileInfo for WSFS. type wsfsFileInfo struct { - oi workspace.ObjectInfo + workspace.ObjectInfo + + // The export format of a notebook. This is not exposed by the SDK. + ReposExportFormat workspace.ExportFormat `json:"repos_export_format,omitempty"` } func (info wsfsFileInfo) Name() string { - return path.Base(info.oi.Path) + return path.Base(info.ObjectInfo.Path) } func (info wsfsFileInfo) Size() int64 { - return info.oi.Size + return info.ObjectInfo.Size } func (info wsfsFileInfo) Mode() fs.FileMode { - switch info.oi.ObjectType { + switch info.ObjectInfo.ObjectType { case workspace.ObjectTypeDirectory, workspace.ObjectTypeRepo: return fs.ModeDir default: @@ -69,7 +73,7 @@ func (info wsfsFileInfo) Mode() fs.FileMode { } func (info wsfsFileInfo) ModTime() time.Time { - return time.UnixMilli(info.oi.ModifiedAt) + return time.UnixMilli(info.ObjectInfo.ModifiedAt) } func (info wsfsFileInfo) IsDir() bool { @@ -77,7 +81,21 @@ func (info wsfsFileInfo) IsDir() bool { } func (info wsfsFileInfo) Sys() any { - return info.oi + return info.ObjectInfo +} + +// UnmarshalJSON is a custom unmarshaller for the wsfsFileInfo struct. +// It must be defined for this type because otherwise the implementation +// of the embedded ObjectInfo type will be used. +func (info *wsfsFileInfo) UnmarshalJSON(b []byte) error { + return marshal.Unmarshal(b, info) +} + +// MarshalJSON is a custom marshaller for the wsfsFileInfo struct. +// It must be defined for this type because otherwise the implementation +// of the embedded ObjectInfo type will be used. +func (info *wsfsFileInfo) MarshalJSON() ([]byte, error) { + return marshal.Marshal(info) } // WorkspaceFilesClient implements the files-in-workspace API. @@ -293,7 +311,22 @@ func (w *WorkspaceFilesClient) Stat(ctx context.Context, name string) (fs.FileIn return nil, err } - info, err := w.workspaceClient.Workspace.GetStatusByPath(ctx, absPath) + var stat wsfsFileInfo + + // Perform bespoke API call because "return_export_info" is not exposed by the SDK. + // We need "repos_export_format" to determine if the file is a py or a ipynb notebook. + // This is not exposed by the SDK so we need to make a direct API call. + err = w.apiClient.Do( + ctx, + http.MethodGet, + "/api/2.0/workspace/get-status", + nil, + map[string]string{ + "path": absPath, + "return_export_info": "true", + }, + &stat, + ) if err != nil { // If we got an API error we deal with it below. var aerr *apierr.APIError @@ -307,5 +340,5 @@ func (w *WorkspaceFilesClient) Stat(ctx context.Context, name string) (fs.FileIn } } - return wsfsFileInfo{*info}, nil + return stat, nil } diff --git a/libs/filer/workspace_files_client_test.go b/libs/filer/workspace_files_client_test.go index 4e9537641f..650b5be682 100644 --- a/libs/filer/workspace_files_client_test.go +++ b/libs/filer/workspace_files_client_test.go @@ -1,8 +1,10 @@ package filer import ( + "encoding/json" "io/fs" "testing" + "time" "github.com/databricks/databricks-sdk-go/service/workspace" "github.com/stretchr/testify/assert" @@ -54,3 +56,40 @@ func TestWorkspaceFilesDirEntry(t *testing.T) { assert.False(t, i1.IsDir()) assert.True(t, i2.IsDir()) } + +func TestWorkspaceFilesClient_wsfsUnmarshal(t *testing.T) { + payload := ` + { + "created_at": 1671030805916, + "language": "PYTHON", + "modified_at": 1671032235392, + "object_id": 795822750063438, + "object_type": "NOTEBOOK", + "path": "/some/path/to/a/notebook", + "repos_export_format": "SOURCE", + "resource_id": "795822750063438" + } + ` + + var info wsfsFileInfo + err := json.Unmarshal([]byte(payload), &info) + require.NoError(t, err) + + // Fields in the object info. + assert.Equal(t, int64(1671030805916), info.CreatedAt) + assert.Equal(t, workspace.LanguagePython, info.Language) + assert.Equal(t, int64(1671032235392), info.ModifiedAt) + assert.Equal(t, int64(795822750063438), info.ObjectId) + assert.Equal(t, workspace.ObjectTypeNotebook, info.ObjectType) + assert.Equal(t, "/some/path/to/a/notebook", info.Path) + assert.Equal(t, workspace.ExportFormatSource, info.ReposExportFormat) + assert.Equal(t, "795822750063438", info.ResourceId) + + // Functions for fs.FileInfo. + assert.Equal(t, "notebook", info.Name()) + assert.Equal(t, int64(0), info.Size()) + assert.Equal(t, fs.ModePerm, info.Mode()) + assert.Equal(t, time.UnixMilli(1671032235392), info.ModTime()) + assert.False(t, info.IsDir()) + assert.NotNil(t, info.Sys()) +} diff --git a/libs/filer/workspace_files_extensions_client.go b/libs/filer/workspace_files_extensions_client.go index 3ce6913af2..a872dcc65a 100644 --- a/libs/filer/workspace_files_extensions_client.go +++ b/libs/filer/workspace_files_extensions_client.go @@ -6,22 +6,17 @@ import ( "fmt" "io" "io/fs" - "net/http" "path" "strings" "github.com/databricks/cli/libs/log" "github.com/databricks/cli/libs/notebook" "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/apierr" - "github.com/databricks/databricks-sdk-go/client" - "github.com/databricks/databricks-sdk-go/marshal" "github.com/databricks/databricks-sdk-go/service/workspace" ) type workspaceFilesExtensionsClient struct { workspaceClient *databricks.WorkspaceClient - apiClient *client.DatabricksClient wsfs Filer root string @@ -35,64 +30,20 @@ var extensionsToLanguages = map[string]workspace.Language{ ".ipynb": workspace.LanguagePython, } -// workspaceFileStatus defines a custom response body for the "/api/2.0/workspace/get-status" API. -// The "repos_export_format" field is not exposed by the SDK. type workspaceFileStatus struct { - *workspace.ObjectInfo - - // The export format of the notebook. This is not exposed by the SDK. - ReposExportFormat workspace.ExportFormat `json:"repos_export_format,omitempty"` + wsfsFileInfo // Name of the file to be used in any API calls made using the workspace files // filer. For notebooks this path does not include the extension. nameForWorkspaceAPI string } -// A custom unmarsaller for the workspaceFileStatus struct. This is needed because -// workspaceFileStatus embeds the workspace.ObjectInfo which itself has a custom -// unmarshaller. -// If a custom unmarshaller is not provided extra fields like ReposExportFormat -// will not have values set. -func (s *workspaceFileStatus) UnmarshalJSON(b []byte) error { - return marshal.Unmarshal(b, s) -} - -func (s *workspaceFileStatus) MarshalJSON() ([]byte, error) { - return marshal.Marshal(s) -} - -func (w *workspaceFilesExtensionsClient) stat(ctx context.Context, name string) (*workspaceFileStatus, error) { - stat := &workspaceFileStatus{ - nameForWorkspaceAPI: name, - } - - // Perform bespoke API call because "return_export_info" is not exposed by the SDK. - // We need "repos_export_format" to determine if the file is a py or a ipynb notebook. - // This is not exposed by the SDK so we need to make a direct API call. - err := w.apiClient.Do( - ctx, - http.MethodGet, - "/api/2.0/workspace/get-status", - nil, - map[string]string{ - "path": path.Join(w.root, name), - "return_export_info": "true", - }, - stat, - ) +func (w *workspaceFilesExtensionsClient) stat(ctx context.Context, name string) (wsfsFileInfo, error) { + info, err := w.wsfs.Stat(ctx, name) if err != nil { - // If we got an API error we deal with it below. - var aerr *apierr.APIError - if !errors.As(err, &aerr) { - return nil, err - } - - // This API returns a 404 if the specified path does not exist. - if aerr.StatusCode == http.StatusNotFound { - return nil, FileDoesNotExistError{path.Join(w.root, name)} - } + return wsfsFileInfo{}, err } - return stat, err + return info.(wsfsFileInfo), err } // This function returns the stat for the provided notebook. The stat object itself contains the path @@ -146,7 +97,10 @@ func (w *workspaceFilesExtensionsClient) getNotebookStatByNameWithExt(ctx contex // Modify the stat object path to include the extension. This stat object will be used // to return the fs.FileInfo object in the stat method. stat.Path = stat.Path + ext - return stat, nil + return &workspaceFileStatus{ + wsfsFileInfo: stat, + nameForWorkspaceAPI: nameWithoutExt, + }, nil } func (w *workspaceFilesExtensionsClient) getNotebookStatByNameWithoutExt(ctx context.Context, name string) (*workspaceFileStatus, error) { @@ -162,7 +116,7 @@ func (w *workspaceFilesExtensionsClient) getNotebookStatByNameWithoutExt(ctx con } // Get the extension for the notebook. - ext := notebook.GetExtensionByLanguage(stat.ObjectInfo) + ext := notebook.GetExtensionByLanguage(&stat.ObjectInfo) // If the notebook was exported as a Jupyter notebook, the extension should be .ipynb. if stat.Language == workspace.LanguagePython && stat.ReposExportFormat == workspace.ExportFormatJupyter { @@ -172,7 +126,10 @@ func (w *workspaceFilesExtensionsClient) getNotebookStatByNameWithoutExt(ctx con // Modify the stat object path to include the extension. This stat object will be used // to return the fs.DirEntry object in the ReadDir method. stat.Path = stat.Path + ext - return stat, nil + return &workspaceFileStatus{ + wsfsFileInfo: stat, + nameForWorkspaceAPI: name, + }, nil } type DuplicatePathError struct { @@ -200,11 +157,6 @@ func (e DuplicatePathError) Error() string { // errors for namespace clashes (e.g. a file and a notebook or a directory and a notebook). // Thus users of these methods should be careful to avoid such clashes. func NewWorkspaceFilesExtensionsClient(w *databricks.WorkspaceClient, root string) (Filer, error) { - apiClient, err := client.New(w.Config) - if err != nil { - return nil, err - } - filer, err := NewWorkspaceFilesClient(w, root) if err != nil { return nil, err @@ -212,7 +164,6 @@ func NewWorkspaceFilesExtensionsClient(w *databricks.WorkspaceClient, root strin return &workspaceFilesExtensionsClient{ workspaceClient: w, - apiClient: apiClient, wsfs: filer, root: root, @@ -240,7 +191,7 @@ func (w *workspaceFilesExtensionsClient) ReadDir(ctx context.Context, name strin return nil, err } // Replace the entry with the new entry that includes the extension. - entries[i] = wsfsDirEntry{wsfsFileInfo{oi: *stat.ObjectInfo}} + entries[i] = wsfsDirEntry{wsfsFileInfo{ObjectInfo: stat.ObjectInfo}} } // Error if we have seen this path before in the current directory. @@ -331,7 +282,7 @@ func (w *workspaceFilesExtensionsClient) Stat(ctx context.Context, name string) return nil, err } - return wsfsFileInfo{oi: *stat.ObjectInfo}, nil + return wsfsFileInfo{ObjectInfo: stat.ObjectInfo}, nil } return info, err