Add custom action with template support

krab/action_custom.go

@@ -3,25 +3,27 @@ package krab
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/ohkrab/krab/cli"
+	"github.com/ohkrab/krab/cliargs"
 	"github.com/ohkrab/krab/krabdb"
+	"github.com/ohkrab/krab/krabtpl"
 	"github.com/ohkrab/krab/tpls"
 )
 
 // ActionCustom keeps data needed to perform this action.
 type ActionCustom struct {
 	Ui         cli.UI
 	Action     *Action
-	Arguments  Arguments
 	Connection krabdb.Connection
 }
 
 func (a *ActionCustom) Help() string {
 	return fmt.Sprint(
 		`Usage: krab action namespace name`,
 		"\n\n",
-		a.Arguments.Help(),
+		a.Action.Arguments.Help(),
 		` 
 Performs custom action.
 `,
@@ -34,10 +36,51 @@ func (a *ActionCustom) Synopsis() string {
 
 // Run in CLI.
 func (a *ActionCustom) Run(args []string) int {
+	ui := a.Ui
+	flags := cliargs.New(args)
+	flags.RequireNonFlagArgs(0)
+
+	for _, arg := range a.Action.Arguments.Args {
+		flags.Add(arg.Name)
+	}
+
+	err := flags.Parse()
+	if err != nil {
+		ui.Output(a.Help())
+		ui.Error(err.Error())
+		return 1
+	}
+
+	err = a.Action.Arguments.Validate(flags.Values())
+	if err != nil {
+		ui.Output(a.Help())
+		ui.Error(err.Error())
+		return 1
+	}
+
+	templates := tpls.New(flags.Values(), krabtpl.Functions)
+
+	err = a.Connection.Get(func(db krabdb.DB) error {
+		return a.Do(context.Background(), db, templates)
+	})
+
+	if err != nil {
+		ui.Error(err.Error())
+		return 1
+	}
+
+	ui.Info("Done")
+
 	return 0
 }
 
 // Do performs the action.
 func (a *ActionCustom) Do(ctx context.Context, db krabdb.DB, tpl *tpls.Templates) error {
-	return nil
+	sb := strings.Builder{}
+	a.Action.ToSQL(&sb)
+	sql := tpl.Render(sb.String())
+
+	_, err := db.ExecContext(ctx, sql)
+
+	return err
 }

krab/action_migrate_down.go

@@ -7,6 +7,7 @@ import (
 	"github.com/ohkrab/krab/cli"
 	"github.com/ohkrab/krab/cliargs"
 	"github.com/ohkrab/krab/krabdb"
+	"github.com/ohkrab/krab/krabtpl"
 	"github.com/ohkrab/krab/tpls"
 	"github.com/pkg/errors"
 )
@@ -73,7 +74,7 @@ func (a *ActionMigrateDown) Run(args []string) int {
 		return 1
 	}
 
-	templates := tpls.New(flags.Values())
+	templates := tpls.New(flags.Values(), krabtpl.Functions)
 
 	a.DownMigration = SchemaMigration{
 		cliargs.Values(flags.Values()).Get("version"),

krab/action_migrate_up.go

@@ -7,6 +7,7 @@ import (
 	"github.com/ohkrab/krab/cli"
 	"github.com/ohkrab/krab/cliargs"
 	"github.com/ohkrab/krab/krabdb"
+	"github.com/ohkrab/krab/krabtpl"
 	"github.com/ohkrab/krab/tpls"
 	"github.com/pkg/errors"
 )
@@ -61,7 +62,7 @@ func (a *ActionMigrateUp) Run(args []string) int {
 		return 1
 	}
 
-	templates := tpls.New(flags.Values())
+	templates := tpls.New(flags.Values(), krabtpl.Functions)
 
 	err = a.Connection.Get(func(db krabdb.DB) error {
 		return a.Do(context.Background(), db, templates, ui)

krab/type_action.go

@@ -10,6 +10,8 @@ type Action struct {
 	Namespace string `hcl:"namespace,label"`
 	RefName   string `hcl:"ref_name,label"`
 
+	Arguments *Arguments `hcl:"arguments,block"`
+
 	SQL string `hcl:"sql"`
 }
 

krabtpl/functions.go

@@ -0,0 +1,11 @@
+package krabtpl
+
+import (
+	"text/template"
+
+	"github.com/ohkrab/krab/krabdb"
+)
+
+var Functions template.FuncMap = template.FuncMap{
+	"quote": krabdb.QuoteIdent,
+}

spec/action_custom_test.go

@@ -38,7 +38,11 @@ migration_set "animals" {
 }
 
 action "view" "refresh" {
-  sql = "REFRESH VIEW anims"
+  arguments {
+    arg "name" {}
+  }
+
+  sql = "REFRESH MATERIALIZED VIEW {{ quote .Args.name }}"
 }
 `))
 	defer c.Teardown()
@@ -48,9 +52,10 @@ action "view" "refresh" {
 
 	_, vals := c.Query(t, "SELECT * FROM anims")
 	if assert.Len(t, vals, 0, "No values should be returned") {
-		c.AssertSuccessfulRun(t, []string{"action", "view", "refresh"})
+		c.AssertSuccessfulRun(t, []string{"action", "view", "refresh", "-name", "anims"})
 		_, vals := c.Query(t, "SELECT * FROM anims")
 		assert.Len(t, vals, 3, "There should be 3 animals after refresh")
+		c.AssertSQLContains(t, `REFRESH MATERIALIZED VIEW "anims"`)
 	}
 
 }

tpls/templates.go

@@ -1,7 +1,7 @@
 package tpls
 
 import (
-	"bytes"
+	"strings"
 	"text/template"
 )
 
@@ -16,10 +16,10 @@ type root struct {
 }
 
 // New created template renderer with values to replace.
-func New(values map[string]interface{}) *Templates {
+func New(values map[string]interface{}, funcMap template.FuncMap) *Templates {
 	t := &Templates{
 		values:   values,
-		template: template.New(""),
+		template: template.New("").Funcs(funcMap),
 	}
 	return t
 }
@@ -32,17 +32,17 @@ func (t *Templates) Validate(s string) error {
 
 // Render applies values and renders final output.
 func (t *Templates) Render(s string) string {
-	w := bytes.NewBufferString("")
+	sb := strings.Builder{}
 	template, err := t.template.Parse(s)
 	if err != nil {
 		//TODO: handle error
 		panic(err)
 	}
 
-	err = template.Execute(w, root{Args: t.values})
+	err = template.Execute(&sb, root{Args: t.values})
 	if err != nil {
 		//TODO: handle error
 		panic(err)
 	}
-	return w.String()
+	return sb.String()
 }