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

     1  package controldisplay
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/csv"
     6  	"fmt"
     7  	"strings"
     8  	"sync"
     9  	"text/template"
    10  	"time"
    11  
    12  	"github.com/Masterminds/sprig/v3"
    13  )
    14  
    15  // templateFuncs merges desired functions from sprig with custom functions that we
    16  // define in steampipe
    17  func templateFuncs(renderContext TemplateRenderContext) template.FuncMap {
    18  	useFromSprigMap := []string{"upper", "toJson", "quote", "dict", "add", "now", "toPrettyJson"}
    19  
    20  	var funcs template.FuncMap = template.FuncMap{}
    21  	sprigMap := sprig.TxtFuncMap()
    22  	for _, use := range useFromSprigMap {
    23  		f, found := sprigMap[use]
    24  		if !found {
    25  			// guarantee that when a function is expected to be present
    26  			// it does not slip through any crack
    27  			panic(fmt.Sprintf("%s not found", use))
    28  		}
    29  		if found {
    30  			funcs[use] = f
    31  		}
    32  	}
    33  	// custom steampipe functions - ones we couldn't find in sprig
    34  	formatterTemplateFuncMap := template.FuncMap{
    35  		"durationInSeconds": durationInSeconds,
    36  		"toCsvCell":         toCSVCellFnFactory(renderContext.Config.Separator),
    37  	}
    38  	for k, v := range formatterTemplateFuncMap {
    39  		funcs[k] = v
    40  	}
    41  
    42  	return funcs
    43  }
    44  
    45  // toCsvCell escapes a value for csv
    46  // we need to do this in a factory function, so that we can
    47  // set the separator for the CSV writer for this render session
    48  func toCSVCellFnFactory(comma string) func(interface{}) string {
    49  	csvWriterBuffer := bytes.NewBufferString("")
    50  	csvBufferLock := sync.Mutex{}
    51  
    52  	csvWriter := csv.NewWriter(csvWriterBuffer)
    53  	if len(comma) > 0 {
    54  		csvWriter.Comma = []rune(comma)[0]
    55  	}
    56  
    57  	return func(v interface{}) string {
    58  		csvBufferLock.Lock()
    59  		defer csvBufferLock.Unlock()
    60  
    61  		csvWriterBuffer.Reset()
    62  		csvWriter.Write([]string{fmt.Sprintf("%v", v)})
    63  		csvWriter.Flush()
    64  		return strings.TrimSpace(csvWriterBuffer.String())
    65  	}
    66  }
    67  
    68  // durationInSeconds returns the passed in duration as seconds
    69  func durationInSeconds(t time.Duration) float64 { return t.Seconds() }