From 4bf32cb666c6cd4c362d905b4d27d0c00875d5be Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Tue, 17 Oct 2023 12:07:55 +0200
Subject: [PATCH] Fix rendering of streaming response (#876)

## Changes

The update to the Go SDK v0.23.0 in #772 included a change to make the
billable usage API return its streaming response. This still did not
make the command print out the CSV returned by the API, however. To do
so, we call `cmdio.RenderReader` in case the response is a byte stream.

Note: there is an opportunity to parse the CSV and return JSON if
requested, but that is out of scope for this PR (it is a rather big
customization of the command).

Fixes #574.

## Tests

Manually confirmed that `databricks account billable-usage download` now
returns CSV.
---
 .codegen/service.go.tmpl                     | 11 ++++++++---
 cmd/account/billable-usage/billable-usage.go |  3 ++-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/.codegen/service.go.tmpl b/.codegen/service.go.tmpl
index 4ede142d..1ee19acb 100644
--- a/.codegen/service.go.tmpl
+++ b/.codegen/service.go.tmpl
@@ -268,9 +268,14 @@ func init() {
 		if err != nil {
 			return err
 		}
-        {{if .Response -}}
-		return cmdio.Render(ctx, response)
-		{{- else -}}
+		{{ if .Response -}}
+			{{- if .IsResponseByteStream -}}
+			defer response.{{.ResponseBodyField.PascalName}}.Close()
+			return cmdio.RenderReader(ctx, response.{{.ResponseBodyField.PascalName}})
+			{{- else -}}
+			return cmdio.Render(ctx, response)
+			{{- end -}}
+		{{ else -}}
 		return nil
 		{{- end -}}
 {{- end -}}
diff --git a/cmd/account/billable-usage/billable-usage.go b/cmd/account/billable-usage/billable-usage.go
index 1dde38f6..5e9b33f8 100755
--- a/cmd/account/billable-usage/billable-usage.go
+++ b/cmd/account/billable-usage/billable-usage.go
@@ -85,7 +85,8 @@ func newDownload() *cobra.Command {
 		if err != nil {
 			return err
 		}
-		return cmdio.Render(ctx, response)
+		defer response.Contents.Close()
+		return cmdio.RenderReader(ctx, response.Contents)
 	}
 
 	// Disable completions since they are not applicable.