package cmdio import ( "bytes" "context" "encoding/base64" "encoding/json" "errors" "fmt" "io" "strings" "text/tabwriter" "text/template" "time" "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go/listing" "github.com/fatih/color" "github.com/nwidger/jsoncolor" ) // Heredoc is the equivalent of compute.TrimLeadingWhitespace // (command-execution API helper from SDK), except it's more // friendly to non-printable characters. func Heredoc(tmpl string) (trimmed string) { lines := strings.Split(tmpl, "\n") leadingWhitespace := 1<<31 - 1 for _, line := range lines { for pos, char := range line { if char == ' ' || char == '\t' { continue } // first non-whitespace character if pos < leadingWhitespace { leadingWhitespace = pos } // is not needed further break } } for i := 0; i < len(lines); i++ { if lines[i] == "" || strings.TrimSpace(lines[i]) == "" { continue } if len(lines[i]) < leadingWhitespace { trimmed += lines[i] + "\n" // or not.. } else { trimmed += lines[i][leadingWhitespace:] + "\n" } } return strings.TrimSpace(trimmed) } // writeFlusher represents a buffered writer that can be flushed. This is useful when // buffering writing a large number of resources (such as during a list API). type writeFlusher interface { io.Writer Flush() error } type jsonRenderer interface { // Render an object as JSON to the provided writeFlusher. renderJson(context.Context, writeFlusher) error } type textRenderer interface { // Render an object as text to the provided writeFlusher. renderText(context.Context, io.Writer) error } type templateRenderer interface { // Render an object using the provided template and write to the provided tabwriter.Writer. renderTemplate(context.Context, *template.Template, *tabwriter.Writer) error } type readerRenderer struct { reader io.Reader } func (r readerRenderer) renderText(_ context.Context, w io.Writer) error { _, err := io.Copy(w, r.reader) return err } type iteratorRenderer[T any] struct { t listing.Iterator[T] bufferSize int } func (ir iteratorRenderer[T]) getBufferSize() int { if ir.bufferSize == 0 { return 20 } return ir.bufferSize } func (ir iteratorRenderer[T]) renderJson(ctx context.Context, w writeFlusher) error { // Iterators are always rendered as a list of resources in JSON. _, err := w.Write([]byte("[\n ")) if err != nil { return err } for i := 0; ir.t.HasNext(ctx); i++ { if i != 0 { _, err = w.Write([]byte(",\n ")) if err != nil { return err } } n, err := ir.t.Next(ctx) if err != nil { return err } res, err := json.MarshalIndent(n, " ", " ") if err != nil { return err } _, err = w.Write(res) if err != nil { return err } if (i+1)%ir.getBufferSize() == 0 { err = w.Flush() if err != nil { return err } } } _, err = w.Write([]byte("\n]\n")) if err != nil { return err } return w.Flush() } func (ir iteratorRenderer[T]) renderTemplate(ctx context.Context, t *template.Template, w *tabwriter.Writer) error { buf := make([]any, 0, ir.getBufferSize()) for i := 0; ir.t.HasNext(ctx); i++ { n, err := ir.t.Next(ctx) if err != nil { return err } buf = append(buf, n) if len(buf) == cap(buf) { err = t.Execute(w, buf) if err != nil { return err } err = w.Flush() if err != nil { return err } buf = buf[:0] } } if len(buf) > 0 { err := t.Execute(w, buf) if err != nil { return err } } return w.Flush() } type defaultRenderer struct { t any } func (d defaultRenderer) renderJson(_ context.Context, w writeFlusher) error { pretty, err := fancyJSON(d.t) if err != nil { return err } _, err = w.Write(pretty) if err != nil { return err } _, err = w.Write([]byte("\n")) if err != nil { return err } return w.Flush() } func (d defaultRenderer) renderTemplate(_ context.Context, t *template.Template, w *tabwriter.Writer) error { return t.Execute(w, d.t) } // Returns something implementing one of the following interfaces: // - jsonRenderer // - textRenderer // - templateRenderer func newRenderer(t any) any { if r, ok := t.(io.Reader); ok { return readerRenderer{reader: r} } return defaultRenderer{t: t} } func newIteratorRenderer[T any](i listing.Iterator[T]) iteratorRenderer[T] { return iteratorRenderer[T]{t: i} } type bufferedFlusher struct { w io.Writer b *bytes.Buffer } func (b bufferedFlusher) Write(bs []byte) (int, error) { return b.b.Write(bs) } func (b bufferedFlusher) Flush() error { _, err := b.w.Write(b.b.Bytes()) if err != nil { return err } b.b.Reset() return nil } func newBufferedFlusher(w io.Writer) writeFlusher { return bufferedFlusher{ w: w, b: &bytes.Buffer{}, } } func renderWithTemplate(r any, ctx context.Context, outputFormat flags.Output, w io.Writer, headerTemplate, template string) error { // TODO: add terminal width & white/dark theme detection switch outputFormat { case flags.OutputJSON: if jr, ok := r.(jsonRenderer); ok { return jr.renderJson(ctx, newBufferedFlusher(w)) } return errors.New("json output not supported") case flags.OutputText: if tr, ok := r.(templateRenderer); ok && template != "" { return renderUsingTemplate(ctx, tr, w, headerTemplate, template) } if tr, ok := r.(textRenderer); ok { return tr.renderText(ctx, w) } if jr, ok := r.(jsonRenderer); ok { return jr.renderJson(ctx, newBufferedFlusher(w)) } return errors.New("no renderer defined") default: return fmt.Errorf("invalid output format: %s", outputFormat) } } type listingInterface interface { HasNext(context.Context) bool } func Render(ctx context.Context, v any) error { c := fromContext(ctx) if _, ok := v.(listingInterface); ok { panic("use RenderIterator instead") } return renderWithTemplate(newRenderer(v), ctx, c.outputFormat, c.out, c.headerTemplate, c.template) } func RenderIterator[T any](ctx context.Context, i listing.Iterator[T]) error { c := fromContext(ctx) return renderWithTemplate(newIteratorRenderer(i), ctx, c.outputFormat, c.out, c.headerTemplate, c.template) } func RenderWithTemplate(ctx context.Context, v any, headerTemplate, template string) error { c := fromContext(ctx) if _, ok := v.(listingInterface); ok { panic("use RenderIteratorWithTemplate instead") } return renderWithTemplate(newRenderer(v), ctx, c.outputFormat, c.out, headerTemplate, template) } func RenderIteratorWithTemplate[T any](ctx context.Context, i listing.Iterator[T], headerTemplate, template string) error { c := fromContext(ctx) return renderWithTemplate(newIteratorRenderer(i), ctx, c.outputFormat, c.out, headerTemplate, template) } func RenderJson(ctx context.Context, v any) error { c := fromContext(ctx) if _, ok := v.(listingInterface); ok { panic("use RenderIteratorJson instead") } return renderWithTemplate(newRenderer(v), ctx, flags.OutputJSON, c.out, c.headerTemplate, c.template) } func RenderIteratorJson[T any](ctx context.Context, i listing.Iterator[T]) error { c := fromContext(ctx) return renderWithTemplate(newIteratorRenderer(i), ctx, c.outputFormat, c.out, c.headerTemplate, c.template) } func renderUsingTemplate(ctx context.Context, r templateRenderer, w io.Writer, headerTmpl, tmpl string) error { tw := tabwriter.NewWriter(w, 0, 4, 2, ' ', 0) base := template.New("command").Funcs(template.FuncMap{ // we render colored output if stdout is TTY, otherwise we render text. // in the future we'll check if we can explicitly check for stderr being // a TTY "header": color.BlueString, "red": color.RedString, "green": color.GreenString, "blue": color.BlueString, "yellow": color.YellowString, "magenta": color.MagentaString, "cyan": color.CyanString, "bold": func(format string, a ...interface{}) string { return color.New(color.Bold).Sprintf(format, a...) }, "italic": func(format string, a ...interface{}) string { return color.New(color.Italic).Sprintf(format, a...) }, "replace": strings.ReplaceAll, "join": strings.Join, "bool": func(v bool) string { if v { return color.GreenString("YES") } return color.RedString("NO") }, "pretty_json": func(in string) (string, error) { var tmp any err := json.Unmarshal([]byte(in), &tmp) if err != nil { return "", err } b, err := fancyJSON(tmp) if err != nil { return "", err } return string(b), nil }, "pretty_date": func(t time.Time) string { return t.Format("2006-01-02T15:04:05Z") }, "b64_encode": func(in string) (string, error) { var out bytes.Buffer enc := base64.NewEncoder(base64.StdEncoding, &out) _, err := enc.Write([]byte(in)) if err != nil { return "", err } err = enc.Close() if err != nil { return "", err } return out.String(), nil }, "b64_decode": func(in string) (string, error) { dec := base64.NewDecoder(base64.StdEncoding, strings.NewReader(in)) out, err := io.ReadAll(dec) if err != nil { return "", err } return string(out), nil }, }) if headerTmpl != "" { headerT, err := base.Parse(headerTmpl) if err != nil { return err } err = headerT.Execute(tw, nil) if err != nil { return err } tw.Write([]byte("\n")) // Do not flush here. Instead, allow the first 100 resources to determine the initial spacing of the header columns. } t, err := base.Parse(tmpl) if err != nil { return err } err = r.renderTemplate(ctx, t, tw) if err != nil { return err } return tw.Flush() } func fancyJSON(v any) ([]byte, error) { // create custom formatter f := jsoncolor.NewFormatter() // set custom colors f.StringColor = color.New(color.FgGreen) f.TrueColor = color.New(color.FgGreen, color.Bold) f.FalseColor = color.New(color.FgRed) f.NumberColor = color.New(color.FgCyan) f.NullColor = color.New(color.FgMagenta) f.ObjectColor = color.New(color.Reset) f.CommaColor = color.New(color.Reset) f.ColonColor = color.New(color.Reset) return jsoncolor.MarshalIndentWithFormatter(v, "", " ", f) }