CLI variables (#2263)

* Add meta-commands \set and \unset (for variables)

* WIP: first pass at variable replacement

* Use a mapReplacer instead of having Command implement replacer

* remove circular reference with variables

* remove the `replacer` interface; just have it be a struct

* use a lexer for variable replacement

cli/cli.go

@@ -79,6 +79,9 @@ type Command struct {
 	// i.e. it will quit after the command is complete.
 	Files []string `json:"files"`
 
+	// variables holds the variables created with the \set meta-command.
+	variables map[string]string
+
 	// nonInteractiveMode is set to true when fbsql is running in
 	// non-ineracative mode. And example of this is when the user has provided a
 	// `-c` flag in the command line.
@@ -89,6 +92,8 @@ type Command struct {
 }
 
 func NewCommand(logdest logger.Logger) *Command {
+	variables := make(map[string]string)
+
 	return &Command{
 		Config: &Config{
 			Host: defaultHost,
@@ -107,8 +112,8 @@ func NewCommand(logdest logger.Logger) *Command {
 			HistoryPath: "",
 		},
 
-		splitter:   newSplitter(),
 		buffer:     newBuffer(),
+		splitter:   newSplitter(newReplacer(variables)),
 		workingDir: newWorkingDir(),
 
 		Stdin:  Stdin,
@@ -118,6 +123,8 @@ func NewCommand(logdest logger.Logger) *Command {
 		output:       Stdout,
 		writeOptions: defaultWriteOptions(),
 
+		variables: variables,
+
 		quit: make(chan struct{}),
 	}
 }

cli/meta.go

@@ -7,6 +7,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"path/filepath"
+	"sort"
 	"strconv"
 	"strings"
 	"syscall"
@@ -54,6 +55,7 @@ var _ metaCommand = (*metaReset)(nil)
 var _ metaCommand = (*metaSet)(nil)
 var _ metaCommand = (*metaTiming)(nil)
 var _ metaCommand = (*metaTuplesOnly)(nil)
+var _ metaCommand = (*metaUnset)(nil)
 var _ metaCommand = (*metaWarn)(nil)
 var _ metaCommand = (*metaWatch)(nil)
 var _ metaCommand = (*metaWrite)(nil)
@@ -333,6 +335,10 @@ Operating System
   \cd [DIR]              change the current working directory
   \timing [on|off]       toggle timing of commands
   \! [COMMAND]           execute command in shell or start interactive shell
+
+Variables
+  \set [NAME [VALUE]]    set internal variable, or list all if no parameters
+  \unset NAME            unset (delete) internal variable
 `
 	cmd.Printf("%s\n", helpText)
 
@@ -367,7 +373,7 @@ func executeFile(cmd *Command, fileName string) (action, error) {
 	}
 	defer file.Close()
 
-	splitter := newSplitter()
+	splitter := newSplitter(newReplacer(cmd.variables))
 	buffer := newBuffer()
 
 	// Read the file by line, pushing the lines into a new line splitter, then
@@ -647,7 +653,29 @@ func newMetaSet(args []string) *metaSet {
 }
 
 func (m *metaSet) execute(cmd *Command) (action, error) {
-	// TODO: set the variable (or clear it, etc)
+	switch len(m.args) {
+	case 0:
+		// Sort the variables before printing them.
+		keys := make([]string, 0, len(cmd.variables))
+		for k := range cmd.variables {
+			keys = append(keys, k)
+		}
+		sort.Strings(keys)
+
+		// Print out the variables.
+		for _, k := range keys {
+			cmd.Printf("%s = '%s'\n", k, cmd.variables[k])
+		}
+		cmd.writeOptions.timing = !cmd.writeOptions.timing
+	default:
+		// The first arg is the key, the remaining args are concatenated together to form the values
+		// For example:
+		// \set one two three
+		// will result in `one = 'twothree'`
+		k := m.args[0]
+		v := strings.Join(m.args[1:], "")
+		cmd.variables[k] = v
+	}
 	return actionNone, nil
 }
 
@@ -729,6 +757,35 @@ func (m *metaTuplesOnly) execute(cmd *Command) (action, error) {
 	return actionNone, nil
 }
 
+// ////////////////////////////////////////////////////////////////////////////
+// unset
+// ////////////////////////////////////////////////////////////////////////////
+type metaUnset struct {
+	args []string
+}
+
+func newMetaUnset(args []string) *metaUnset {
+	return &metaUnset{
+		args: args,
+	}
+}
+
+func (m *metaUnset) execute(cmd *Command) (action, error) {
+	switch len(m.args) {
+	case 0:
+		cmd.Printf("\\unset: missing required argument\n")
+		return actionNone, nil
+	default:
+		if len(m.args) > 1 {
+			for _, s := range m.args[1:] {
+				cmd.Printf("\\unset: extra argument \"%s\" ignored\n", s)
+			}
+		}
+		delete(cmd.variables, m.args[0])
+	}
+	return actionNone, nil
+}
+
 // ////////////////////////////////////////////////////////////////////////////
 // warn
 // ////////////////////////////////////////////////////////////////////////////
@@ -854,7 +911,7 @@ func (m *metaWrite) execute(cmd *Command) (action, error) {
 //	`cmd 'arg1' arg2 'arg three'`
 //
 // It returns the metaCommand which maps to `cmd`.
-func splitMetaCommand(in string) (metaCommand, error) {
+func splitMetaCommand(in string, replacer *replacer) (metaCommand, error) {
 	parts := strings.SplitN(in, ` `, 2)
 	key := strings.TrimRightFunc(parts[0], unicode.IsSpace)
 
@@ -877,6 +934,11 @@ func splitMetaCommand(in string) (metaCommand, error) {
 		}
 	}
 
+	// Do variable replacement.
+	for i := range args {
+		args[i] = replacer.replace(args[i])
+	}
+
 	switch key {
 	case "!":
 		return newMetaBang(args), nil
@@ -916,6 +978,8 @@ func splitMetaCommand(in string) (metaCommand, error) {
 		return newMetaTuplesOnly(args), nil
 	case "timing":
 		return newMetaTiming(args), nil
+	case "unset":
+		return newMetaUnset(args), nil
 	case "warn":
 		return newMetaWarn(args), nil
 	case "watch":

cli/replacer.go

@@ -0,0 +1,109 @@
+package cli
+
+import (
+	"strings"
+
+	"github.com/benhoyt/goawk/lexer"
+)
+
+// replacer can replace parts of a string based on some rules and the provided
+// map[string]string. For example, the Command can replace strings with values
+// in its `variables` map.
+type replacer struct {
+	m map[string]string
+}
+
+func newReplacer(m map[string]string) *replacer {
+	return &replacer{
+		m: m,
+	}
+}
+
+// replace replaces all instances of the string pattern `:key` with the value at
+// m[key]. For example we want something like this:
+//
+// GIVEN: `start :one,:'two', :"three" ::four ::`
+//
+// with map
+//
+//	map[string]string{
+//		"one":   "repl1",
+//		"three": "repl3",
+//	}
+//
+// WANT: `start repl1,:'two', "repl3" ::four ::`
+func (r *replacer) replace(s string) string {
+	// If no variables have been added to the map, there's no need to parse the
+	// string for variable replacement.
+	if len(r.m) == 0 {
+		return s
+	}
+
+	line := []byte(s)
+	lex := lexer.NewLexer(line)
+
+	// finger contains the index into line at the start of non-variable text
+	// that we want to include, as-is in the output.
+	var finger int
+
+	// sb builds the string which will be the final output.
+	var sb strings.Builder
+	for {
+		pos, tok, _ := lex.Scan()
+
+		switch tok {
+		case lexer.COLON:
+			// last is the last normal character position before the colon.
+			last := pos.Column - 1
+
+			// Get the next byte to see if the colon value is quoted, and if so,
+			// whether its has single or double quotes.
+			b := lex.PeekByte()
+
+			// padding is the amount of padding we have to consider around the
+			// variable name. If the variable is not quoted, it doesn't require
+			// any padding. But if it has quotes, it needs 2 characters of
+			// paddings to accomodate the quotes.
+			padding := 0
+
+			// quote holds the character to use to quote the final, replaced
+			// output value. Because the lexer doesn't tell us how a certain
+			// `string` token was quoted, we need to keep track of that here so
+			// we can put them back.
+			quote := ""
+			switch b {
+			case byte('\''): // single quote
+				quote = `'`
+				padding = 2
+			case byte('"'): // double quote
+				quote = `"`
+				padding = 2
+			}
+
+			pos, tok, key := lex.Scan()
+			switch tok {
+			case lexer.NAME, lexer.STRING:
+				// Write the normal text up to the variable replacement
+				// position.
+				sb.Write(line[finger:last])
+
+				if v, ok := r.m[key]; ok {
+					// Write replaced variable with the quotes it had.
+					sb.WriteString(quote + v + quote)
+				} else {
+					// Since the variable was not found in the map, just write
+					// back what was already there.
+					sb.WriteString(":" + quote + key + quote)
+				}
+
+				// Reset finger to point to the next position after the
+				// variable.
+				finger = pos.Column + len(key) + padding - 1
+			}
+		case lexer.EOF:
+			// Write the remainder of the string and return.
+			sb.Write(line[finger:])
+			return sb.String()
+		}
+	}
+}

cli/replacer_test.go

@@ -0,0 +1,109 @@
+package cli
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestReplacer(t *testing.T) {
+	t.Run("general replace function", func(t *testing.T) {
+
+		m := map[string]string{
+			"v1": "newVone",
+			"v2": "newVtwo",
+		}
+
+		tests := []struct {
+			s   string
+			m   map[string]string
+			exp string
+		}{
+			{
+				// no variables present
+				s:   "foo",
+				m:   m,
+				exp: "foo",
+			},
+			{
+				// variable prefix, but not in map
+				s:   ":foo",
+				m:   m,
+				exp: ":foo",
+			},
+			{
+				// variable name match, but missing prefix
+				s:   "v1",
+				m:   m,
+				exp: "v1",
+			},
+			{
+				// variable name match
+				s:   ":v1",
+				m:   m,
+				exp: "newVone",
+			},
+			{
+				// two variables, the same, no space
+				s:   ":v1:v1",
+				m:   m,
+				exp: "newVonenewVone",
+			},
+			{
+				// two variables, different, no space
+				s:   ":v1:v2",
+				m:   m,
+				exp: "newVonenewVtwo",
+			},
+			{
+				// two variables, different, spaces
+				s:   ":v1  :v2",
+				m:   m,
+				exp: "newVone  newVtwo",
+			},
+			{
+				// one variable, one non-variable, no space
+				s:   ":v1:foo",
+				m:   m,
+				exp: "newVone:foo",
+			},
+			{
+				// one non-variable, one variable, no space
+				s:   "foo:v1",
+				m:   m,
+				exp: "foonewVone",
+			},
+			{
+				// two variables, different, comma
+				s:   ":v1, :v2",
+				m:   m,
+				exp: "newVone, newVtwo",
+			},
+			{
+				// single quotes
+				s:   ":'v1'",
+				m:   m,
+				exp: "'newVone'",
+			},
+			{
+				// double quotes
+				s:   `:"v2"`,
+				m:   m,
+				exp: `"newVtwo"`,
+			},
+			{
+				// more quotes
+				s:   `start :v1,:'two', :"v2" ::four :: `,
+				m:   m,
+				exp: `start newVone,:'two', "newVtwo" ::four :: `,
+			},
+		}
+		for i, test := range tests {
+			t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
+				replacer := newReplacer(test.m)
+				assert.Equal(t, test.exp, replacer.replace(test.s))
+			})
+		}
+	})
+}

cli/splitter.go

@@ -10,10 +10,14 @@ import (
 // metaCommands. It may not be necessary to have this be a separate struct since
 // it contains no members and just has the one `split()` method, but here we
 // are.
-type splitter struct{}
+type splitter struct {
+	replacer *replacer
+}
 
-func newSplitter() *splitter {
-	return &splitter{}
+func newSplitter(r *replacer) *splitter {
+	return &splitter{
+		replacer: r,
+	}
 }
 
 // split splits the given line into queryParts and metaCommands.
@@ -65,6 +69,11 @@ func (s *splitter) splitQueryParts(line string) ([]queryPart, error) {
 	// Look for a termination character;
 	parts := strings.Split(line, terminationChar)
 
+	// Do variable replacement.
+	for i := range parts {
+		parts[i] = s.replacer.replace(parts[i])
+	}
+
 	if len(parts) == 1 {
 		part0 := strings.TrimSpace(parts[0])
 		return []queryPart{
@@ -95,7 +104,7 @@ func (s *splitter) splitQueryParts(line string) ([]queryPart, error) {
 func (s *splitter) splitMetaCommands(in string) ([]metaCommand, error) {
 	parts := strings.Split(in, `\`)
 	if len(parts) == 1 {
-		mc, err := splitMetaCommand(parts[0])
+		mc, err := splitMetaCommand(parts[0], s.replacer)
 		if err != nil {
 			return nil, errors.Wrapf(err, "splitting meta command: %s", parts[0])
 		}
@@ -108,7 +117,7 @@ func (s *splitter) splitMetaCommands(in string) ([]metaCommand, error) {
 		if part == "" {
 			continue
 		}
-		mc, err := splitMetaCommand(part)
+		mc, err := splitMetaCommand(part, s.replacer)
 		if err != nil {
 			return nil, errors.Wrapf(err, "splitting meta command: %s", part)
 		}

cli/splitter_test.go

@@ -8,7 +8,7 @@ import (
 )
 
 func TestSplitter(t *testing.T) {
-	s := newSplitter()
+	s := newSplitter(newReplacer(nil))
 	t.Run("Split", func(t *testing.T) {
 		tests := []struct {
 			line            string

go.mod

@@ -83,6 +83,7 @@ require (
 	github.com/PaesslerAG/gval v1.0.0
 	github.com/PaesslerAG/jsonpath v0.1.1
 	github.com/apache/arrow/go/v10 v10.0.0-20221021053532-2f627c213fc3
+	github.com/benhoyt/goawk v1.21.0
 	github.com/gomem/gomem v0.1.0
 	github.com/google/uuid v1.3.0
 	github.com/jaffee/commandeer v0.6.0

go.sum

@@ -139,6 +139,8 @@ github.com/beevik/ntp v0.3.0 h1:xzVrPrE4ziasFXgBVBZJDP0Wg/KpMwk2KHJ4Ba8GrDw=
 github.com/beevik/ntp v0.3.0/go.mod h1:hIHWr+l3+/clUnF44zdK+CWW7fO8dR5cIylAQ76NRpg=
 github.com/benbjohnson/immutable v0.4.0 h1:CTqXbEerYso8YzVPxmWxh2gnoRQbbB9X1quUC8+vGZA=
 github.com/benbjohnson/immutable v0.4.0/go.mod h1:iAr8OjJGLnLmVUr9MZ/rz4PWUy6Ouc2JLYuMArmvAJM=
+github.com/benhoyt/goawk v1.21.0 h1:GASuhJXHMFZ/2TJBPh+2Ah3kclVGNvGjt+uh3ajMdLk=
+github.com/benhoyt/goawk v1.21.0/go.mod h1:UG1Ld6CjkkHhoyQmErQGSTwmavsTqFnCDYsLSJbovqU=
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=