From a42b8299684a35a11de08e64afb446dec4bb6e0e Mon Sep 17 00:00:00 2001
From: Ilya Kuznetsov <ilya.kuznetsov@databricks.com>
Date: Tue, 12 Nov 2024 13:41:39 +0100
Subject: [PATCH] feat: Support notebooks without extensions

---
 bundle/config/mutator/translate_paths.go        | 17 +++++++++++++++++
 libs/filer/workspace_files_client.go            |  7 +++++++
 libs/filer/workspace_files_extensions_client.go |  6 ++++--
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/bundle/config/mutator/translate_paths.go b/bundle/config/mutator/translate_paths.go
index 321fa5b3..e16bbba3 100644
--- a/bundle/config/mutator/translate_paths.go
+++ b/bundle/config/mutator/translate_paths.go
@@ -11,6 +11,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/notebook"
@@ -117,6 +118,10 @@ func (t *translateContext) rewritePath(
 	return nil
 }
 
+type FileInfoWithVirtualExtension interface {
+	VirtualExtension() string
+}
+
 func (t *translateContext) translateNotebookPath(literal, localFullPath, localRelPath, remotePath string) (string, error) {
 	nb, _, err := notebook.DetectWithFS(t.b.SyncRoot, filepath.ToSlash(localRelPath))
 	if errors.Is(err, fs.ErrNotExist) {
@@ -129,6 +134,18 @@ func (t *translateContext) translateNotebookPath(literal, localFullPath, localRe
 		return "", ErrIsNotNotebook{localFullPath}
 	}
 
+	if config.IsExplicitlyEnabled(t.b.Config.Presets.InPlaceDeployment) {
+		stat, err := t.b.SyncRoot.Stat(filepath.ToSlash(localRelPath))
+		if err != nil {
+			return "", fmt.Errorf("unable to get file info for %s: %w", localFullPath, err)
+		}
+
+		if i, ok := stat.(FileInfoWithVirtualExtension); ok {
+			vext := i.VirtualExtension()
+			return strings.TrimSuffix(remotePath, vext), nil
+		}
+	}
+
 	// Upon import, notebooks are stripped of their extension.
 	return strings.TrimSuffix(remotePath, filepath.Ext(localFullPath)), nil
 }
diff --git a/libs/filer/workspace_files_client.go b/libs/filer/workspace_files_client.go
index 4bb03aea..e4f16a64 100644
--- a/libs/filer/workspace_files_client.go
+++ b/libs/filer/workspace_files_client.go
@@ -53,6 +53,9 @@ type wsfsFileInfo struct {
 
 	// The export format of a notebook. This is not exposed by the SDK.
 	ReposExportFormat workspace.ExportFormat `json:"repos_export_format,omitempty"`
+
+	// Explicitly added file extension of the notebook based on its language
+	VirtualFileExtension string `json:"virtual_file_extension,omitempty"`
 }
 
 func (info wsfsFileInfo) Name() string {
@@ -88,6 +91,10 @@ func (info wsfsFileInfo) WorkspaceObjectInfo() workspace.ObjectInfo {
 	return info.ObjectInfo
 }
 
+func (info wsfsFileInfo) VirtualExtension() string {
+	return info.VirtualFileExtension
+}
+
 // 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.
diff --git a/libs/filer/workspace_files_extensions_client.go b/libs/filer/workspace_files_extensions_client.go
index b24ecf7e..ca340366 100644
--- a/libs/filer/workspace_files_extensions_client.go
+++ b/libs/filer/workspace_files_extensions_client.go
@@ -98,6 +98,7 @@ 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
+	stat.VirtualFileExtension = ext
 	return &workspaceFileStatus{
 		wsfsFileInfo:        stat,
 		nameForWorkspaceAPI: nameWithoutExt,
@@ -127,6 +128,7 @@ 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
+	stat.VirtualFileExtension = ext
 	return &workspaceFileStatus{
 		wsfsFileInfo:        stat,
 		nameForWorkspaceAPI: name,
@@ -214,7 +216,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{ObjectInfo: stat.ObjectInfo}}
+			entries[i] = wsfsDirEntry{wsfsFileInfo{ObjectInfo: stat.ObjectInfo, VirtualFileExtension: stat.VirtualFileExtension}}
 		}
 
 		// Error if we have seen this path before in the current directory.
@@ -313,7 +315,7 @@ func (w *workspaceFilesExtensionsClient) Stat(ctx context.Context, name string)
 			return nil, err
 		}
 
-		return wsfsFileInfo{ObjectInfo: stat.ObjectInfo}, nil
+		return wsfsFileInfo{ObjectInfo: stat.ObjectInfo, VirtualFileExtension: stat.VirtualFileExtension}, nil
 	}
 
 	return info, err