Add ability to template YAML playbook with custom functions (closes #167

)

sql_runner/db.go

@@ -14,30 +14,7 @@ package main
 
 import (
 	"bytes"
-	"math/rand"
-	"os"
-	"strconv"
 	"text/template"
-	"time"
-)
-
-var (
-	templFuncs = template.FuncMap{
-		"nowWithFormat": func(format string) string {
-			return time.Now().Format(format)
-		},
-		"systemEnv": func(env string) string {
-			return os.Getenv(env)
-		},
-		"randomInt": func() (string, error) {
-			r := rand.NewSource(time.Now().UnixNano())
-			return strconv.FormatInt(r.Int63(), 10), nil
-		},
-		"awsChainCredentials":   awsChainCredentials,
-		"awsEC2RoleCredentials": awsEC2RoleCredentials,
-		"awsEnvCredentials":     awsEnvCredentials,
-		"awsProfileCredentials": awsProfileCredentials,
-	}
 )
 
 // Generalized interface to a database client
@@ -70,8 +47,7 @@ func prepareQuery(queryPath string, sp SQLProvider, template bool, variables map
 
 // Fills in a script which is a template
 func fillTemplate(script string, variables map[string]interface{}) (string, error) {
-
-	t, err := template.New("playbook").Funcs(templFuncs).Parse(script)
+	t, err := template.New("playbook").Funcs(TemplFuncs).Parse(script)
 	if err != nil {
 		return "", err
 	}

sql_runner/template.go

@@ -0,0 +1,40 @@
+//
+// Copyright (c) 2015-2020 Snowplow Analytics Ltd. All rights reserved.
+//
+// This program is licensed to you under the Apache License Version 2.0,
+// and you may not use this file except in compliance with the Apache License Version 2.0.
+// You may obtain a copy of the Apache License Version 2.0 at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the Apache License Version 2.0 is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the Apache License Version 2.0 for the specific language governing permissions and limitations there under.
+//
+package main
+
+import (
+	"math/rand"
+	"os"
+	"strconv"
+	"text/template"
+	"time"
+)
+
+var (
+	TemplFuncs = template.FuncMap{
+		"nowWithFormat": func(format string) string {
+			return time.Now().Format(format)
+		},
+		"systemEnv": func(env string) string {
+			return os.Getenv(env)
+		},
+		"randomInt": func() (string, error) {
+			r := rand.NewSource(time.Now().UnixNano())
+			return strconv.FormatInt(r.Int63(), 10), nil
+		},
+		"awsChainCredentials":   awsChainCredentials,
+		"awsEC2RoleCredentials": awsEC2RoleCredentials,
+		"awsEnvCredentials":     awsEnvCredentials,
+		"awsProfileCredentials": awsProfileCredentials,
+	}
+)

sql_runner/yaml_utils.go

@@ -59,7 +59,7 @@ func cleanYaml(rawYaml []byte) []byte {
 }
 
 func fillPlaybookTemplate(playbookStr string, variables map[string]string) (string, error) {
-	t, err := template.New("playbook").Parse(playbookStr)
+	t, err := template.New("playbook").Funcs(TemplFuncs).Parse(playbookStr)
 	if err != nil {
 		return "", err
 	}