implement EXIT command

.vscode/launch.json

@@ -24,7 +24,7 @@
             "request": "launch",
             "mode" : "auto",
             "program": "${workspaceFolder}/cmd/sqlcmd",
-            "args" : ["-Q", "\"select 100 as Count\""],
+            "args" : ["-Q", "EXIT(select 100 as Count)"],
         },
     ]
 }

README.md

@@ -18,6 +18,7 @@ We will be implementing as many command line switches and behaviors as possible
 
 - `-R` switch will be removed. The go runtime does not provide access to user locale information, and it's not readily available through syscall on all supported platforms.
 - Some behaviors that were kept to maintain compatibility with `OSQL` may be changed, such as alignment of column headers for some data types.
+- All commands must fit on one line, even `EXIT`. Interactive mode will not check for open parentheses or quotes for commands and prompt for successive lines. The native sqlcmd allows the query run by `EXIT(query)` to span multiple lines.
 
 ### Packages
 

pkg/sqlcmd/batch_test.go

@@ -11,7 +11,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestBatchNextReset(t *testing.T) {
+func TestBatchNext(t *testing.T) {
 	tests := []struct {
 		s     string
 		stmts []string
@@ -30,6 +30,9 @@ func TestBatchNextReset(t *testing.T) {
 		{"$(x) $(y) 100\nquit", []string{"$(x) $(y) 100"}, []string{"QUIT"}, "-"},
 		{"select 1\n:list", []string{"select 1"}, []string{"LIST"}, "-"},
 		{"select 1\n:reset", []string{"select 1"}, []string{"RESET"}, "-"},
+		{"select 1\n:exit()", []string{"select 1"}, []string{"EXIT"}, "-"},
+		{"select 1\n:exit (select 10)", []string{"select 1"}, []string{"EXIT"}, "-"},
+		{"select 1\n:exit", []string{"select 1"}, []string{"EXIT"}, "-"},
 	}
 	for _, test := range tests {
 		b := NewBatch(sp(test.s, "\n"), newCommands())
@@ -48,7 +51,6 @@ func TestBatchNextReset(t *testing.T) {
 			case err != nil:
 				t.Fatalf("test %s did not expect error, got: %v", test.s, err)
 			}
-			// resetting the buffer for every command purely for test purposes
 			if cmd != nil {
 				cmds = append(cmds, cmd.name)
 			}

pkg/sqlcmd/commands.go

@@ -30,6 +30,11 @@ type Commands map[string]*Command
 func newCommands() Commands {
 	// Commands is the set of Command implementations
 	return map[string]*Command{
+		"EXIT": {
+			regex:  regexp.MustCompile(`(?im)^[\t ]*?:?EXIT(?:[ \t]*(\(?.*\)?$)|$)`),
+			action: exitCommand,
+			name:   "EXIT",
+		},
 		"QUIT": {
 			regex:  regexp.MustCompile(`(?im)^[\t ]*?:?QUIT(?:[ \t]+(.*$)|$)`),
 			action: quitCommand,
@@ -105,6 +110,35 @@ func (c Commands) SetBatchTerminator(terminator string) error {
 	return nil
 }
 
+// exitCommand has 3 modes.
+// With no (), it just exits without running any query
+// With () it runs whatever batch is in the buffer then exits
+// With any text between () it runs the text as a query then exits
+func exitCommand(s *Sqlcmd, args []string, line uint) error {
+	if len(args) == 0 {
+		return ErrExitRequested
+	}
+	params := strings.TrimSpace(args[0])
+	if params == "" {
+		return ErrExitRequested
+	}
+	if !strings.HasPrefix(params, "(") || !strings.HasSuffix(params, ")") {
+		return InvalidCommandError("EXIT", line)
+	}
+	// First we run the current batch
+	query := s.batch.String()
+	if query != "" {
+		query = s.getRunnableQuery(query)
+		_ = s.runQuery(query)
+	}
+	query = strings.TrimSpace(params[1 : len(params)-1])
+	if query != "" {
+		query = s.getRunnableQuery(query)
+		s.Exitcode = s.runQuery(query)
+	}
+	return ErrExitRequested
+}
+
 // quitCommand immediately exits the program without running any more batches
 func quitCommand(s *Sqlcmd, args []string, line uint) error {
 	if args != nil && strings.TrimSpace(args[0]) != "" {
@@ -132,38 +166,8 @@ func goCommand(s *Sqlcmd, args []string, line uint) error {
 		return nil
 	}
 	query = s.getRunnableQuery(query)
-	// This loop will likely be refactored to a helper when we implement -Q and :EXIT(query)
 	for i := 0; i < n; i++ {
-
-		s.Format.BeginBatch(query, s.vars, s.GetOutput(), s.GetError())
-		rows, qe := s.db.Query(query)
-		if qe != nil {
-			s.Format.AddError(qe)
-		}
-
-		results := true
-		for qe == nil && results {
-			cols, err := rows.ColumnTypes()
-			if err != nil {
-				s.Format.AddError(err)
-			} else {
-				s.Format.BeginResultSet(cols)
-				active := rows.Next()
-				for active {
-					s.Format.AddRow(rows)
-					active = rows.Next()
-				}
-				if err = rows.Err(); err != nil {
-					s.Format.AddError(err)
-				}
-				s.Format.EndResultSet()
-			}
-			results = rows.NextResultSet()
-			if err = rows.Err(); err != nil {
-				s.Format.AddError(err)
-			}
-		}
-		s.Format.EndBatch()
+		_ = s.runQuery(query)
 	}
 	s.batch.Reset(nil)
 	return nil

pkg/sqlcmd/commands_test.go

@@ -38,6 +38,9 @@ func TestCommandParsing(t *testing.T) {
 		{` :Error c:\folder\file`, "ERROR", []string{`c:\folder\file`}},
 		{`:Setvar A1 "some value" `, "SETVAR", []string{`A1 "some value" `}},
 		{` :Listvar`, "LISTVAR", []string{""}},
+		{`:EXIT (select 100 as count)`, "EXIT", []string{"(select 100 as count)"}},
+		{`:EXIT ( )`, "EXIT", []string{"( )"}},
+		{`EXIT `, "EXIT", []string{""}},
 	}
 
 	for _, test := range commands {

pkg/sqlcmd/format.go

@@ -28,8 +28,8 @@ type Formatter interface {
 	BeginResultSet([]*sql.ColumnType)
 	// EndResultSet is called after all rows in a result set have been processed
 	EndResultSet()
-	// AddRow is called for each row in a result set
-	AddRow(*sql.Rows)
+	// AddRow is called for each row in a result set. It returns the value of the first column
+	AddRow(*sql.Rows) string
 	// AddMessage is called for every information message returned by the server during the batch
 	AddMessage(string)
 	// AddError is called for each error encountered during batch execution
@@ -137,19 +137,20 @@ func (f *sqlCmdFormatterType) EndResultSet() {
 }
 
 // Writes the current row to the designated output writer
-func (f *sqlCmdFormatterType) AddRow(row *sql.Rows) {
-
-	f.writepos = 0
+func (f *sqlCmdFormatterType) AddRow(row *sql.Rows) string {
+	retval := ""
 	values, err := f.scanRow(row)
 	if err != nil {
 		f.mustWriteErr(err.Error())
-		return
+		return retval
 	}
 
 	// values are the full values, look at the displaywidth of each column and truncate accordingly
 	for i, v := range values {
 		if i > 0 {
 			f.writeOut(f.vars.ColumnSeparator())
+		} else {
+			retval = v
 		}
 		f.printColumnValue(v, i)
 	}
@@ -160,6 +161,8 @@ func (f *sqlCmdFormatterType) AddRow(row *sql.Rows) {
 		f.printColumnHeadings()
 	}
 	f.writeOut(SqlcmdEol)
+	return retval
+
 }
 
 // Writes a non-error message to the designated message writer

pkg/sqlcmd/sqlcmd.go

@@ -98,11 +98,13 @@ func (s *Sqlcmd) Run(once bool, processAll bool) error {
 		var args []string
 		var err error
 		if s.Query != "" {
-			cmd = s.Cmd["GO"]
-			args = make([]string, 0)
 			s.batch.Reset([]rune(s.Query))
 			// batch.Next validates variable syntax
-			_, _, err = s.batch.Next()
+			cmd, args, err = s.batch.Next()
+			if cmd == nil {
+				cmd = s.Cmd["GO"]
+				args = make([]string, 0)
+			}
 			s.Query = ""
 		} else {
 			cmd, args, err = s.batch.Next()
@@ -381,3 +383,57 @@ func setupCloseHandler(s *Sqlcmd) {
 		os.Exit(0)
 	}()
 }
+
+// runQuery runs the query and prints the results
+// The return value is based on the first cell of the last column of the last result set.
+// If it's numeric, it will be converted to int
+// -100 : Error encountered prior to selecting return value
+// -101: No rows found
+// -102: Conversion error occurred when selecting return value
+func (s *Sqlcmd) runQuery(query string) int {
+	retcode := -101
+	s.Format.BeginBatch(query, s.vars, s.GetOutput(), s.GetError())
+	rows, qe := s.db.Query(query)
+	if qe != nil {
+		s.Format.AddError(qe)
+	}
+	var err error
+	var cols []*sql.ColumnType
+	results := true
+	for qe == nil && results {
+		cols, err = rows.ColumnTypes()
+		if err != nil {
+			retcode = -100
+			s.Format.AddError(err)
+		} else {
+			s.Format.BeginResultSet(cols)
+			active := rows.Next()
+			for active {
+				col1 := s.Format.AddRow(rows)
+				active = rows.Next()
+				if !active {
+					if col1 == "" {
+						retcode = 0
+					} else if _, cerr := fmt.Sscanf(col1, "%d", &retcode); cerr != nil {
+						retcode = -102
+					}
+				}
+			}
+
+			if retcode != -102 {
+				if err = rows.Err(); err != nil {
+					retcode = -100
+					s.Format.AddError(err)
+				}
+			}
+			s.Format.EndResultSet()
+		}
+		results = rows.NextResultSet()
+		if err = rows.Err(); err != nil {
+			retcode = -100
+			s.Format.AddError(err)
+		}
+	}
+	s.Format.EndBatch()
+	return retcode
+}

pkg/sqlcmd/sqlcmd_test.go

@@ -4,6 +4,7 @@
 package sqlcmd
 
 import (
+	"bytes"
 	"database/sql"
 	"fmt"
 	"os"
@@ -207,6 +208,32 @@ func TestGetRunnableQuery(t *testing.T) {
 
 }
 
+func TestExitInitialQuery(t *testing.T) {
+	s, buf := setupSqlCmdWithMemoryOutput(t)
+	s.Query = "EXIT(SELECT '1200', 2100)"
+	err := s.Run(true, false)
+	if assert.NoError(t, err, "s.Run(once = true)") {
+		s.SetOutput(nil)
+		o := buf.buf.String()
+		assert.Equal(t, "1200 2100"+SqlcmdEol+SqlcmdEol, o, "Output")
+		assert.Equal(t, 1200, s.Exitcode, "ExitCode")
+	}
+
+}
+
+func setupSqlCmdWithMemoryOutput(t testing.TB) (*Sqlcmd, *memoryBuffer) {
+	v := InitializeVariables(true)
+	v.Set(SQLCMDMAXVARTYPEWIDTH, "0")
+	s := New(nil, "", v)
+	s.Connect.Password = os.Getenv(SQLCMDPASSWORD)
+	s.Format = NewSQLCmdDefaultFormatter(true)
+	buf := &memoryBuffer{buf: new(bytes.Buffer)}
+	s.SetOutput(buf)
+	err := s.ConnectDb("", "", "", true)
+	assert.NoError(t, err, "s.ConnectDB")
+	return s, buf
+}
+
 func setupSqlcmdWithFileOutput(t testing.TB) (*Sqlcmd, *os.File) {
 	v := InitializeVariables(true)
 	v.Set(SQLCMDMAXVARTYPEWIDTH, "0")