Add support to passing variables as command-line options

integration-tests/postgres-sql/good/assert.sql

@@ -7,10 +7,10 @@ DECLARE
   expected_average_age CONSTANT integer := 25;
 BEGIN
   {{/* Update view names to today's date or Travis will error */}}
-  IF (SELECT average_age <> expected_average_age FROM {{.test_schema}}.view_2015_09_13) THEN
+  IF (SELECT average_age <> expected_average_age FROM {{.test_schema}}.view_2015_09_14) THEN
     RAISE EXCEPTION 'Average_age % does not match expected age %',
-    	(SELECT average_age FROM {{.test_schema}}.view_2015_09_13),
-    	expected_average_age;
+      (SELECT average_age FROM {{.test_schema}}.view_2015_09_14),
+      expected_average_age;
   END IF;
 END;
 $$ LANGUAGE plpgsql;

main.go

@@ -35,7 +35,7 @@ func main() {
 
 	options := processFlags()
 
-	pb, err := playbook.ParsePlaybook(options.playbook)
+	pb, err := playbook.ParsePlaybook(options.playbook, options.variables)
 	if err != nil {
 		log.Fatalf("Could not parse playbook (YAML): %s", err.Error())
 	}
@@ -51,7 +51,7 @@ func main() {
 // required flag
 func processFlags() Options {
 
-	var options Options
+	var options Options = NewOptions()
 	var fs = options.GetFlagSet()
 	fs.Parse(os.Args[1:])
 

options.go

@@ -15,13 +15,36 @@ package main
 import (
 	"flag"
 	"fmt"
+	"strings"
 )
 
+type CLIVariables map[string]string
+
+// Implement the Value interface
+func (i *CLIVariables) String() string {
+	return fmt.Sprintf("%s", *i)
+}
+
+func (i *CLIVariables) Set(value string) error {
+	var split = strings.SplitN(value, "=", 2)
+	if len(split) > 1 {
+		key := split[0]
+		val := split[1]
+		(*i)[key] = val
+	}
+	return nil
+}
+
 type Options struct {
-	help     bool
-	version  bool
-	playbook string
-	sqlroot  string
+	help      bool
+	version   bool
+	playbook  string
+	sqlroot   string
+	variables CLIVariables
+}
+
+func NewOptions() Options {
+	return Options{variables: make(map[string]string)}
 }
 
 func (o *Options) GetFlagSet() *flag.FlagSet {
@@ -31,6 +54,7 @@ func (o *Options) GetFlagSet() *flag.FlagSet {
 	fs.BoolVar(&(o.version), "version", false, "Shows the program version")
 	fs.StringVar(&(o.playbook), "playbook", "", "Playbook of SQL scripts to execute")
 	fs.StringVar(&(o.sqlroot), "sqlroot", SQLROOT_PLAYBOOK, fmt.Sprintf("Absolute path to SQL scripts. Use %s and %s for those respective paths", SQLROOT_PLAYBOOK, SQLROOT_BINARY))
+	fs.Var(&(o.variables), "var", "Variables to be passed to the playbook, in the key=value format")
 	// TODO: add format flag if/when we support TOML
 
 	return fs

playbook/playbook.go

@@ -35,7 +35,18 @@ type Query struct {
 }
 
 // Dispatch to format-specific parser
-func ParsePlaybook(filepath string) (Playbook, error) {
+func ParsePlaybook(filepath string, variables map[string]string) (Playbook, error) {
 	// TODO: Add TOML support?
-	return parsePlaybookYaml(filepath)
+	playbook, err := parsePlaybookYaml(filepath)
+	if err == nil {
+		playbook = MergeCLIVariables(playbook, variables)
+	}
+	return playbook, err
+}
+
+func MergeCLIVariables(playbook Playbook, variables map[string]string) Playbook {
+	for k, v := range variables {
+		playbook.Variables[k] = v
+	}
+	return playbook
 }