github.com/turbot/steampipe@v1.7.0-rc.0.0.20240517123944-7cef272d4458/pkg/control/controldisplay/formatter_template.go (about)

     1  package controldisplay
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"text/template"
     9  
    10  	"github.com/spf13/viper"
    11  	"github.com/turbot/steampipe/pkg/constants"
    12  	"github.com/turbot/steampipe/pkg/control/controlexecute"
    13  	"github.com/turbot/steampipe/pkg/utils"
    14  	"github.com/turbot/steampipe/pkg/version"
    15  )
    16  
    17  // TemplateFormatter implements the 'Formatter' interface and exposes a generic template based output mechanism
    18  // for 'check' execution trees
    19  type TemplateFormatter struct {
    20  	template     *template.Template
    21  	exportFormat *OutputTemplate
    22  }
    23  
    24  func NewTemplateFormatter(input *OutputTemplate) (*TemplateFormatter, error) {
    25  	templateFuncs := templateFuncs(TemplateRenderContext{})
    26  
    27  	// add a stub "render_context" function
    28  	// this will be overwritten before we execute the template
    29  	// if we don't put this here, then templates which use this
    30  	// won't parse and will throw Error: template: ****: function "render_context" not defined
    31  	templateFuncs["render_context"] = func() TemplateRenderContext { return TemplateRenderContext{} }
    32  
    33  	t, err := template.New("outlet").
    34  		Funcs(templateFuncs).
    35  		ParseFS(os.DirFS(input.TemplatePath), "*")
    36  
    37  	if err != nil {
    38  		return nil, fmt.Errorf("could not load template '%s' - %v", input.TemplatePath, err)
    39  	}
    40  
    41  	return &TemplateFormatter{exportFormat: input, template: t}, nil
    42  }
    43  
    44  func (tf TemplateFormatter) Format(ctx context.Context, tree *controlexecute.ExecutionTree) (io.Reader, error) {
    45  	reader, writer := io.Pipe()
    46  	go func() {
    47  		workingDirectory, err := os.Getwd()
    48  		if err != nil {
    49  			writer.CloseWithError(err)
    50  			return
    51  		}
    52  		renderContext := TemplateRenderContext{
    53  			Constants: TemplateRenderConstants{
    54  				SteampipeVersion: version.SteampipeVersion.String(),
    55  				WorkingDir:       workingDirectory,
    56  			},
    57  			Config: TemplateRenderConfig{
    58  				RenderHeader: viper.GetBool(constants.ArgHeader),
    59  				Separator:    viper.GetString(constants.ArgSeparator),
    60  			},
    61  			Data: tree,
    62  		}
    63  
    64  		// overwrite the "render_context" function to return the current render context
    65  		templateFuncs := templateFuncs(renderContext)
    66  		templateFuncs["render_context"] = func() TemplateRenderContext { return renderContext }
    67  
    68  		t, err := tf.template.Clone()
    69  		if err != nil {
    70  			writer.CloseWithError(err)
    71  			return
    72  		}
    73  		t = t.Funcs(templateFuncs)
    74  
    75  		if err := t.ExecuteTemplate(writer, "output", renderContext); err != nil {
    76  			writer.CloseWithError(err)
    77  		} else {
    78  			writer.Close()
    79  		}
    80  	}()
    81  
    82  	// tactical - for json, prettify the output
    83  	if tf.shouldPrettify() {
    84  		return utils.PrettifyJsonFromReader(reader)
    85  	}
    86  
    87  	return reader, nil
    88  }
    89  
    90  func (tf TemplateFormatter) FileExtension() string {
    91  	return tf.exportFormat.FileExtension
    92  }
    93  
    94  func (tf TemplateFormatter) Name() string {
    95  	return tf.exportFormat.FormatName
    96  }
    97  
    98  func (tf TemplateFormatter) Alias() string {
    99  	if tf.exportFormat.FormatFullName != tf.exportFormat.FormatName {
   100  		return tf.exportFormat.FormatFullName
   101  	}
   102  	return ""
   103  }
   104  
   105  func (tf TemplateFormatter) shouldPrettify() bool {
   106  	return tf.Name() == constants.OutputFormatJSON
   107  }