Skip to content

Commit

Permalink
Support text/plain response as streaming request body (#638)
Browse files Browse the repository at this point in the history
## Changes
For APIs using the `text/plain` content type (today, only usage
download), support these by simply streaming the response body back to
the end user. This way, they can read the response and save it to a file
or post-process it further.

## Tests
Made a simple app that downloads billable usage to a temporary file:
```go
package main

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"os"

	databricks "github.com/databricks/databricks-sdk-go"
	"github.com/databricks/databricks-sdk-go/logger"
	"github.com/databricks/databricks-sdk-go/service/billing"
)

func writeToTempFile(r io.ReadCloser) (*os.File, error) {
	// Ensure r is closed after reading
	defer r.Close()

	// Create a temporary file
	tmpFile, err := ioutil.TempFile("", "prefix-")
	if err != nil {
		return nil, err
	}

	// Copy data from io.ReadCloser to the temporary file
	_, err = io.Copy(tmpFile, r)
	if err != nil {
		tmpFile.Close() // Close and remove temporary file in case of errors
		os.Remove(tmpFile.Name())
		return nil, err
	}

	// Rewind the temporary file for further reading
	tmpFile.Seek(0, io.SeekStart)

	return tmpFile, nil
}

func main() {
	logger.DefaultLogger = &logger.SimpleLogger{Level: logger.LevelTrace}
	a, err := databricks.NewAccountClient(&databricks.Config{
		Profile:      "...",
		DebugHeaders: true,
	})
	if err != nil {
		panic(err)
	}
	ctx := context.Background()
	res, err := a.BillableUsage.Download(ctx, billing.DownloadRequest{
		StartMonth:   "2023-08",
		EndMonth:     "2023-09",
		PersonalData: true,
	})
	if err != nil {
		panic(err)
	}
	tmp, err := writeToTempFile(res.Contents)
	if err != nil {
		panic(err)
	}
	fmt.Printf("tmp: %v\n", tmp.Name())
}

```
This works.

Note that this test depends on a change to the content type for the
OpenAPI spec for usage download to match the true content type.
  • Loading branch information
mgyucht authored Oct 9, 2023
1 parent 410c4b0 commit d821075
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions openapi/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,7 @@ type MimeType string
const (
MimeTypeJson MimeType = "application/json"
MimeTypeOctetStream MimeType = "application/octet-stream"
MimeTypeTextPlain MimeType = "text/plain"
)

// IsByteStream returns true if the body should be modeled as a byte stream.
Expand All @@ -311,6 +312,7 @@ func (m MimeType) IsByteStream() bool {
var allowedMimeTypes = []MimeType{
MimeTypeJson,
MimeTypeOctetStream,
MimeTypeTextPlain,
}

func (b *Body) MimeTypeAndMediaType() (MimeType, *MediaType) {
Expand Down

0 comments on commit d821075

Please sign in to comment.