diff --git a/cli/cli.go b/cli/cli.go index 7a3ee1606..cc8f2e0a3 100644 --- a/cli/cli.go +++ b/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{}), } } diff --git a/cli/meta.go b/cli/meta.go index b249691fd..142f9b78a 100644 --- a/cli/meta.go +++ b/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": diff --git a/cli/replacer.go b/cli/replacer.go new file mode 100644 index 000000000..b22330a8a --- /dev/null +++ b/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() + } + } +} diff --git a/cli/replacer_test.go b/cli/replacer_test.go new file mode 100644 index 000000000..4cf1069ce --- /dev/null +++ b/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)) + }) + } + }) +} diff --git a/cli/splitter.go b/cli/splitter.go index ca62eb257..bb3c5f2d1 100644 --- a/cli/splitter.go +++ b/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) } diff --git a/cli/splitter_test.go b/cli/splitter_test.go index 54b0fb3eb..ae22251f5 100644 --- a/cli/splitter_test.go +++ b/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 diff --git a/go.mod b/go.mod index f92aa51ab..d4bd82bed 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index fdbd41d61..ba7e26a9f 100644 --- a/go.sum +++ b/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=