Added Upload and Download methods where applicable

.codegen/workspaces.go.tmpl

@@ -11,6 +11,8 @@ import (
 type WorkspaceClient struct {
 	Config *config.Config
 
+	Files *files.FilesAPI
+
 	{{range .Services}}{{if not .IsAccounts}}
 	{{.Comment "    // " 80}}
 	{{.Name}} *{{.Package.Name}}.{{.Name}}API
@@ -34,6 +36,7 @@ func NewWorkspaceClient(c ...*Config) (*WorkspaceClient, error) {
 	}
 	return &WorkspaceClient{
 		Config: cfg,
+		Files: files.NewFiles(apiClient),
 		{{range .Services}}{{if not .IsAccounts}}
 		{{.Name}}: {{.Package.Name}}.New{{.Name}}(apiClient),
 		{{- end}}{{end}}

.gitignore

@@ -25,3 +25,4 @@ raw-specs/**
 out/**
 *.log
 .databricks/
+.DS_Store

client/client.go

@@ -101,6 +101,11 @@ func (c *DatabricksClient) unmarshal(body []byte, response any) error {
 	if len(body) == 0 {
 		return nil
 	}
+	// If the destination is bytes.Buffer, write the body over there
+	if raw, ok := response.(*bytes.Buffer); ok {
+		_, err := raw.Write(body)
+		return err
+	}
 	// If the destination is a byte slice, pass the body verbatim.
 	if raw, ok := response.(*[]byte); ok {
 		*raw = body
@@ -113,7 +118,6 @@ func (c *DatabricksClient) addHostToRequestUrl(r *http.Request) error {
 	if r.URL == nil {
 		return fmt.Errorf("no URL found in request")
 	}
-	r.Header.Set("Content-Type", "application/json")
 	url, err := url.Parse(c.Config.Host)
 	if err != nil {
 		return err
@@ -123,6 +127,11 @@ func (c *DatabricksClient) addHostToRequestUrl(r *http.Request) error {
 	return nil
 }
 
+func (c *DatabricksClient) addApplicationJsonContentType(r *http.Request) error {
+	r.Header.Set("Content-Type", "application/json")
+	return nil
+}
+
 func (c *DatabricksClient) redactedDump(prefix string, body []byte) (res string) {
 	return bodyLogger{
 		debugTruncateBytes: c.debugTruncateBytes,
@@ -275,6 +284,7 @@ func (c *DatabricksClient) perform(
 	visitors = append([]func(*http.Request) error{
 		c.Config.Authenticate,
 		c.addHostToRequestUrl,
+		c.addApplicationJsonContentType,
 		c.addAuthHeaderToUserAgent,
 	}, visitors...)
 	resp, err := retries.Poll(ctx, c.retryTimeout,

internal/dbfs_test.go → internal/files_test.go

@@ -8,6 +8,7 @@ import (
 	"math/rand"
 	"path"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -24,6 +25,22 @@ func (buf hashable) Hash() uint32 {
 	return h.Sum32()
 }
 
+func TestAccFilesAPI(t *testing.T) {
+	t.SkipNow() // until available on prod
+	ctx, w := workspaceTest(t)
+
+	filePath := RandomName("/Volumes/bogdanghita/default/v3_shared/sdk-testing/txt-")
+	err := w.Files.Upload(ctx, filePath, strings.NewReader("abcd"))
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		err = w.Files.Delete(ctx, filePath)
+		assert.NoError(t, err)
+	})
+	raw, err := w.Files.ReadFile(ctx, filePath)
+	require.NoError(t, err)
+	assert.Equal(t, "abcd", string(raw))
+}
+
 func TestAccDbfsOpen(t *testing.T) {
 	ctx, w := workspaceTest(t)
 	if w.Config.IsGcp() {

internal/workspace_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/databricks/databricks-sdk-go"
@@ -53,9 +54,8 @@ func TestAccWorkspaceIntegration(t *testing.T) {
 
 	// Export the notebook and assert the contents
 	exportResponse, err := w.Workspace.Export(ctx, workspace.ExportRequest{
-		DirectDownload: false,
-		Format:         workspace.ExportFormatSource,
-		Path:           notebook,
+		Format: workspace.ExportFormatSource,
+		Path:   notebook,
 	})
 	require.NoError(t, err)
 	assert.True(t, exportResponse.Content == base64.StdEncoding.EncodeToString([]byte("# Databricks notebook source\nprint('hello from job')")))
@@ -74,18 +74,92 @@ func TestAccWorkspaceIntegration(t *testing.T) {
 	assert.Contains(t, paths, notebook)
 }
 
+func TestAccWorkspaceUploadNotebookWithFileExtensionNoTranspile(t *testing.T) {
+	// TODO: remove NoTranspile suffix once other languages get Upload/Donwload features
+	ctx, w := workspaceTest(t)
+
+	notebookPath := filepath.Join("/Users", me(t, w).UserName, RandomName("notebook-")+".py")
+
+	err := w.Workspace.Upload(ctx, notebookPath, strings.NewReader("print(1)"))
+	assert.NoError(t, err)
+	t.Cleanup(func() {
+		err = w.Workspace.Delete(ctx, workspace.Delete{
+			Path: notebookPath,
+		})
+		require.NoError(t, err)
+	})
+
+	info, err := w.Workspace.GetStatusByPath(ctx, notebookPath)
+	assert.NoError(t, err)
+	assert.Equal(t, workspace.LanguagePython, info.Language)
+	assert.Equal(t, workspace.ObjectTypeNotebook, info.ObjectType)
+
+	contents, err := w.Workspace.ReadFile(ctx, notebookPath)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "# Databricks notebook source\nprint(1)", string(contents))
+}
+
+func TestAccWorkspaceUploadNotebookWithFileNoExtensionNoTranspile(t *testing.T) {
+	// TODO: remove NoTranspile suffix once other languages get Upload/Donwload features
+	ctx, w := workspaceTest(t)
+
+	notebookPath := filepath.Join("/Users", me(t, w).UserName, RandomName("notebook-"))
+
+	err := w.Workspace.Upload(ctx, notebookPath, strings.NewReader("print(1)"),
+		workspace.UploadLanguage(workspace.LanguagePython))
+	assert.NoError(t, err)
+	t.Cleanup(func() {
+		err = w.Workspace.Delete(ctx, workspace.Delete{
+			Path: notebookPath,
+		})
+		require.NoError(t, err)
+	})
+
+	info, err := w.Workspace.GetStatusByPath(ctx, notebookPath)
+	assert.NoError(t, err)
+	assert.Equal(t, workspace.LanguagePython, info.Language)
+	assert.Equal(t, workspace.ObjectTypeNotebook, info.ObjectType)
+
+	contents, err := w.Workspace.ReadFile(ctx, notebookPath)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "# Databricks notebook source\nprint(1)", string(contents))
+}
+
+func TestAccWorkspaceUploadFileNoTranspile(t *testing.T) {
+	// TODO: remove NoTranspile suffix once other languages get Upload/Donwload features
+	ctx, w := workspaceTest(t)
+
+	txtPath := filepath.Join("/Users", me(t, w).UserName, RandomName("txt-"))
+
+	err := w.Workspace.Upload(ctx, txtPath, strings.NewReader("print(1)"),
+		workspace.UploadFormat(workspace.ExportFormatAuto))
+	assert.NoError(t, err)
+	t.Cleanup(func() {
+		err = w.Workspace.Delete(ctx, workspace.Delete{
+			Path: txtPath,
+		})
+		require.NoError(t, err)
+	})
+
+	info, err := w.Workspace.GetStatusByPath(ctx, txtPath)
+	assert.NoError(t, err)
+	assert.Equal(t, workspace.ObjectTypeFile, info.ObjectType)
+
+	contents, err := w.Workspace.ReadFile(ctx, txtPath)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "print(1)", string(contents))
+}
+
 func TestAccWorkspaceRecursiveListNoTranspile(t *testing.T) {
 	ctx, w := workspaceTest(t)
 	notebook := myNotebookPath(t, w)
 
 	// Import the test notebook
-	err := w.Workspace.Import(ctx, workspace.Import{
-		Path:      notebook,
-		Format:    workspace.ExportFormatSource,
-		Language:  workspace.LanguagePython,
-		Content:   base64.StdEncoding.EncodeToString([]byte("# Databricks notebook source\nprint('hello from job')")),
-		Overwrite: true,
-	})
+	err := w.Workspace.Upload(ctx, notebook, strings.NewReader("print(1)"),
+		workspace.UploadOverwrite())
 	require.NoError(t, err)
 
 	allMyNotebooks, err := w.Workspace.RecursiveList(ctx, filepath.Join("/Users", me(t, w).UserName))

openapi/code/package.go

@@ -385,6 +385,11 @@ func (pkg *Package) Load(spec *openapi.Specification, tag *openapi.Tag) error {
 					if seenParams[param.Name] {
 						continue
 					}
+					if prefix == "/api/2.0/workspace/export" && param.Name == "direct_download" {
+						// prevent changing the response content type via request parameter
+						// https://github.com/databricks/databricks-sdk-py/issues/104
+						continue
+					}
 					params = append(params, *param)
 					seenParams[param.Name] = true
 				}

openapi/roll/tool.go

@@ -31,7 +31,7 @@ func NewSuite(dirname string) (*suite, error) {
 			// not transpilable
 			return nil
 		}
-		if strings.HasSuffix(path, "dbfs_test.go") {
+		if strings.HasSuffix(path, "files_test.go") {
 			// not transpilable
 			return nil
 		}

service/files/files_api.go

@@ -0,0 +1,81 @@
+package files
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+
+	"github.com/databricks/databricks-sdk-go/client"
+)
+
+// BEGIN TODO: remove this once FilesAPI gets OpenAPI spec
+
+func NewFiles(client *client.DatabricksClient) *FilesAPI {
+	return &FilesAPI{
+		impl: &filesImpl{
+			client: client,
+		},
+	}
+}
+
+type FilesService interface{}
+
+type FilesAPI struct {
+	impl FilesService
+}
+
+type filesImpl struct {
+	client *client.DatabricksClient
+}
+
+// END TODO: remove this once FilesAPI gets OpenAPI spec
+
+func (a *FilesAPI) Upload(ctx context.Context, path string, r io.Reader) error {
+	impl, ok := a.impl.(*filesImpl)
+	if !ok {
+		return fmt.Errorf("wrong impl: %v", a.impl)
+	}
+	return impl.client.Do(ctx, "PUT", "/api/2.0/fs/files"+path, r, nil,
+		func(r *http.Request) error {
+			r.Header.Set("Content-Type", "application/octet-stream")
+			return nil
+		})
+}
+
+// WriteFile is identical to [os.WriteFile] but for Files API.
+func (a *FilesAPI) WriteFile(ctx context.Context, name string, data []byte) error {
+	return a.Upload(ctx, name, bytes.NewBuffer(data))
+}
+
+func (a *FilesAPI) Download(ctx context.Context, path string) (io.ReadCloser, error) {
+	impl, ok := a.impl.(*filesImpl)
+	if !ok {
+		return nil, fmt.Errorf("wrong impl: %v", a.impl)
+	}
+	var buf bytes.Buffer
+	err := impl.client.Do(ctx, "GET", "/api/2.0/fs/files"+path, nil, &buf)
+	if err != nil {
+		return nil, err
+	}
+	// TODO: direclty call lower-level APIs to return raw HTTP response stream
+	return io.NopCloser(&buf), nil
+}
+
+// ReadFile is identical to [os.ReadFile] but for Files API.
+func (a *FilesAPI) ReadFile(ctx context.Context, name string) ([]byte, error) {
+	b, err := a.Download(ctx, name)
+	if err != nil {
+		return nil, err
+	}
+	return io.ReadAll(b)
+}
+
+func (a *FilesAPI) Delete(ctx context.Context, path string) error {
+	impl, ok := a.impl.(*filesImpl)
+	if !ok {
+		return fmt.Errorf("wrong impl: %v", a.impl)
+	}
+	return impl.client.Do(ctx, "DELETE", "/api/2.0/fs/files"+path, nil, nil)
+}