implement error handling switches (#49)

Author: shueybubbles

cmd/sqlcmd/main.go

@@ -42,6 +42,10 @@ type SQLCmdArguments struct {
 	WorkstationName             string            `short:"H" help:"This option sets the sqlcmd scripting variable SQLCMDWORKSTATION. The workstation name is listed in the hostname column of the sys.sysprocesses catalog view and can be returned using the stored procedure sp_who. If this option is not specified, the default is the current computer name. This name can be used to identify different sqlcmd sessions."`
 	ApplicationIntent           string            `short:"K" default:"default" enum:"default,ReadOnly" help:"Declares the application workload type when connecting to a server. The only currently supported value is ReadOnly. If -K is not specified, the sqlcmd utility will not support connectivity to a secondary replica in an Always On availability group."`
 	EncryptConnection           string            `short:"N" default:"default" enum:"default,false,true,disable" help:"This switch is used by the client to request an encrypted connection."`
+	DriverLoggingLevel          int               `help:"Level of mssql driver messages to print."`
+	ExitOnError                 bool              `short:"b" help:"Specifies that sqlcmd exits and returns a DOS ERRORLEVEL value when an error occurs."`
+	ErrorSeverityLevel          uint8             `short:"V" help:"Controls the severity level that is used to set the ERRORLEVEL variable on exit."`
+	ErrorLevel                  int               `short:"m" help:"Controls which error messages are sent to stdout. Messages that have severity level greater than or equal to this level are sent."`
 }
 
 // Validate accounts for settings not described by Kong attributes
@@ -122,7 +126,7 @@ func setVars(vars *sqlcmd.Variables, args *SQLCmdArguments) {
 		},
 		sqlcmd.SQLCMDWORKSTATION: func(a *SQLCmdArguments) string { return args.WorkstationName },
 		sqlcmd.SQLCMDSERVER:      func(a *SQLCmdArguments) string { return a.Server },
-		sqlcmd.SQLCMDERRORLEVEL:  func(a *SQLCmdArguments) string { return "" },
+		sqlcmd.SQLCMDERRORLEVEL:  func(a *SQLCmdArguments) string { return fmt.Sprint(a.ErrorLevel) },
 		sqlcmd.SQLCMDPACKETSIZE: func(a *SQLCmdArguments) string {
 			if args.PacketSize > 0 {
 				return fmt.Sprint(args.PacketSize)
@@ -165,6 +169,9 @@ func setConnect(s *sqlcmd.Sqlcmd, args *SQLCmdArguments) {
 	s.Connect.Encrypt = args.EncryptConnection
 	s.Connect.PacketSize = args.PacketSize
 	s.Connect.WorkstationName = args.WorkstationName
+	s.Connect.LogLevel = args.DriverLoggingLevel
+	s.Connect.ExitOnError = args.ExitOnError
+	s.Connect.ErrorSeverityLevel = args.ErrorSeverityLevel
 }
 
 func run(vars *sqlcmd.Variables) (int, error) {

cmd/sqlcmd/main_test.go

@@ -67,6 +67,9 @@ func TestValidCommandLineToArgsConversion(t *testing.T) {
 		{[]string{"-a", "550", "-l", "45", "-H", "mystation", "-K", "ReadOnly", "-N", "true"}, func(args SQLCmdArguments) bool {
 			return args.PacketSize == 550 && args.LoginTimeout == 45 && args.WorkstationName == "mystation" && args.ApplicationIntent == "ReadOnly" && args.EncryptConnection == "true"
 		}},
+		{[]string{"-b", "-m", "15", "-V", "20"}, func(args SQLCmdArguments) bool {
+			return args.ExitOnError && args.ErrorLevel == 15 && args.ErrorSeverityLevel == 20
+		}},
 	}
 
 	for _, test := range commands {

pkg/sqlcmd/commands.go

@@ -129,12 +129,15 @@ func exitCommand(s *Sqlcmd, args []string, line uint) error {
 	query := s.batch.String()
 	if query != "" {
 		query = s.getRunnableQuery(query)
-		_ = s.runQuery(query)
+		if exitCode, err := s.runQuery(query); err != nil {
+			s.Exitcode = exitCode
+			return ErrExitRequested
+		}
 	}
 	query = strings.TrimSpace(params[1 : len(params)-1])
 	if query != "" {
 		query = s.getRunnableQuery(query)
-		s.Exitcode = s.runQuery(query)
+		s.Exitcode, _ = s.runQuery(query)
 	}
 	return ErrExitRequested
 }
@@ -167,7 +170,10 @@ func goCommand(s *Sqlcmd, args []string, line uint) error {
 	}
 	query = s.getRunnableQuery(query)
 	for i := 0; i < n; i++ {
-		_ = s.runQuery(query)
+		if retcode, err := s.runQuery(query); err != nil {
+			s.Exitcode = retcode
+			return err
+		}
 	}
 	s.batch.Reset(nil)
 	return nil

pkg/sqlcmd/format.go

@@ -172,16 +172,21 @@ func (f *sqlCmdFormatterType) AddMessage(msg string) {
 
 // Writes an error to the designated err Writer
 func (f *sqlCmdFormatterType) AddError(err error) {
+	print := true
 	b := new(strings.Builder)
 	msg := err.Error()
 	switch e := (err).(type) {
 	case mssql.Error:
-		b.WriteString(fmt.Sprintf("Msg %d, Level %d, State %d, Server %s, Line %d%s", e.Number, e.Class, e.State, e.ServerName, e.LineNo, SqlcmdEol))
-		msg = strings.TrimPrefix(msg, "mssql: ")
+		if print = f.vars.ErrorLevel() <= 0 || e.Class >= uint8(f.vars.ErrorLevel()); print {
+			b.WriteString(fmt.Sprintf("Msg %d, Level %d, State %d, Server %s, Line %d%s", e.Number, e.Class, e.State, e.ServerName, e.LineNo, SqlcmdEol))
+			msg = strings.TrimPrefix(msg, "mssql: ")
+		}
+	}
+	if print {
+		b.WriteString(msg)
+		b.WriteString(SqlcmdEol)
+		f.mustWriteOut(fitToScreen(b, f.vars.ScreenWidth()).String())
 	}
-	b.WriteString(msg)
-	b.WriteString(SqlcmdEol)
-	f.mustWriteErr(fitToScreen(b, f.vars.ScreenWidth()).String())
 }
 
 // Prints column headings based on columnDetail, variables, and command line arguments

pkg/sqlcmd/sqlcmd.go

@@ -20,6 +20,7 @@ import (
 	"syscall"
 
 	mssql "github.com/denisenkom/go-mssqldb"
+	"github.com/denisenkom/go-mssqldb/msdsn"
 	"github.com/gohxs/readline"
 	"github.com/golang-sql/sqlexp"
 )
@@ -57,6 +58,10 @@ type ConnectSettings struct {
 	WorkstationName string
 	// ApplicationIntent can only be empty or "ReadOnly"
 	ApplicationIntent string
+	// mssql driver log level
+	LogLevel           int
+	ExitOnError        bool
+	ErrorSeverityLevel uint8
 }
 
 func (c ConnectSettings) authenticationMethod() string {
@@ -96,6 +101,7 @@ func New(l *readline.Instance, workingDirectory string, vars *Variables) *Sqlcmd
 		Cmd:              newCommands(),
 	}
 	s.batch = NewBatch(s.scanNext, s.Cmd)
+	mssql.SetContextLogger(s)
 	return s
 }
 
@@ -157,9 +163,17 @@ func (s *Sqlcmd) Run(once bool, processAll bool) error {
 			if err != nil {
 				fmt.Fprintln(stderr, err)
 				lastError = err
-				continue
 			}
 		}
+		if err != nil && s.Connect.ExitOnError {
+			// If the error were due to a SQL error, the GO command handler
+			// would have set ExitCode already
+			if s.Exitcode == 0 {
+				s.Exitcode = 1
+			}
+			lastError = err
+			break
+		}
 		if execute {
 			s.Query = s.batch.String()
 			once = true
@@ -260,6 +274,9 @@ func (s *Sqlcmd) ConnectionString() (connectionString string, err error) {
 	if s.Connect.Encrypt != "" && s.Connect.Encrypt != "default" {
 		query.Add("encrypt", s.Connect.Encrypt)
 	}
+	if s.Connect.LogLevel > 0 {
+		query.Add("log", fmt.Sprint(s.Connect.LogLevel))
+	}
 	connectionURL.RawQuery = query.Encode()
 	return connectionURL.String(), nil
 }
@@ -449,7 +466,7 @@ func (s *Sqlcmd) sqlAuthentication() bool {
 // -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 {
+func (s *Sqlcmd) runQuery(query string) (int, error) {
 	retcode := -101
 	s.Format.BeginBatch(query, s.vars, s.GetOutput(), s.GetError())
 	ctx := context.Background()
@@ -469,6 +486,7 @@ func (s *Sqlcmd) runQuery(query string) int {
 			s.Format.AddMessage(m.Message)
 		case sqlexp.MsgError:
 			s.Format.AddError(m.Error)
+			qe = s.handleError(&retcode, m.Error)
 		case sqlexp.MsgRowsAffected:
 			if m.Count == 1 {
 				s.Format.AddMessage("(1 row affected)")
@@ -479,6 +497,7 @@ func (s *Sqlcmd) runQuery(query string) int {
 			results = rows.NextResultSet()
 			if err = rows.Err(); err != nil {
 				retcode = -100
+				qe = s.handleError(&retcode, err)
 				s.Format.AddError(err)
 			}
 			if results {
@@ -492,6 +511,7 @@ func (s *Sqlcmd) runQuery(query string) int {
 					cols, err = rows.ColumnTypes()
 					if err != nil {
 						retcode = -100
+						qe = s.handleError(&retcode, err)
 						s.Format.AddError(err)
 					} else {
 						s.Format.BeginResultSet(cols)
@@ -510,12 +530,51 @@ func (s *Sqlcmd) runQuery(query string) int {
 			if retcode != -102 {
 				if err = rows.Err(); err != nil {
 					retcode = -100
+					qe = s.handleError(&retcode, err)
 					s.Format.AddError(err)
 				}
 			}
 			s.Format.EndResultSet()
 		}
 	}
 	s.Format.EndBatch()
-	return retcode
+	return retcode, qe
+}
+
+// returns ErrExitRequested if the error is a SQL error and satisfies the connection's error handling configuration
+func (s *Sqlcmd) handleError(retcode *int, err error) error {
+	if err == nil {
+		return nil
+	}
+
+	var minSeverityToExit uint8 = 11
+	if s.Connect.ErrorSeverityLevel > 0 {
+		minSeverityToExit = s.Connect.ErrorSeverityLevel
+	}
+	var errSeverity uint8
+	switch sqlError := err.(type) {
+	case mssql.Error:
+		errSeverity = sqlError.Class
+	}
+
+	if s.Connect.ErrorSeverityLevel > 0 {
+		if errSeverity >= minSeverityToExit {
+			*retcode = int(errSeverity)
+			s.Exitcode = *retcode
+		}
+	} else if s.Connect.ExitOnError {
+		if errSeverity >= minSeverityToExit {
+			*retcode = 1
+		}
+	}
+	if s.Connect.ExitOnError && errSeverity >= minSeverityToExit {
+		return ErrExitRequested
+	}
+	return nil
+}
+
+// Log attempts to write driver traces to the current output. It ignores errors
+func (s Sqlcmd) Log(_ context.Context, _ msdsn.Log, msg string) {
+	_, _ = s.GetOutput().Write([]byte("DRIVER:" + msg))
+	_, _ = s.GetOutput().Write([]byte(SqlcmdEol))
 }

pkg/sqlcmd/sqlcmd_test.go

@@ -7,6 +7,7 @@ import (
 	"bytes"
 	"database/sql"
 	"fmt"
+	"io"
 	"os"
 	"os/user"
 	"strings"
@@ -240,6 +241,76 @@ func TestExitInitialQuery(t *testing.T) {
 
 }
 
+func TestExitCodeSetOnError(t *testing.T) {
+	s, _ := setupSqlCmdWithMemoryOutput(t)
+	s.Connect.ErrorSeverityLevel = 12
+	retcode, err := s.runQuery("RAISERROR (N'Testing!' , 11, 1)")
+	assert.NoError(t, err, "!ExitOnError 11")
+	assert.Equal(t, -101, retcode, "Raiserror below ErrorSeverityLevel")
+	retcode, err = s.runQuery("RAISERROR (N'Testing!' , 14, 1)")
+	assert.NoError(t, err, "!ExitOnError 14")
+	assert.Equal(t, 14, retcode, "Raiserror above ErrorSeverityLevel")
+	s.Connect.ExitOnError = true
+	retcode, err = s.runQuery("RAISERROR (N'Testing!' , 11, 1)")
+	assert.NoError(t, err, "ExitOnError and Raiserror below ErrorSeverityLevel")
+	assert.Equal(t, -101, retcode, "Raiserror below ErrorSeverityLevel")
+	retcode, err = s.runQuery("RAISERROR (N'Testing!' , 14, 1)")
+	assert.ErrorIs(t, err, ErrExitRequested, "ExitOnError and Raiserror above ErrorSeverityLevel")
+	assert.Equal(t, 14, retcode, "ExitOnError and Raiserror above ErrorSeverityLevel")
+	s.Connect.ErrorSeverityLevel = 0
+	retcode, err = s.runQuery("RAISERROR (N'Testing!' , 11, 1)")
+	assert.ErrorIs(t, err, ErrExitRequested, "ExitOnError and ErrorSeverityLevel = 0, Raiserror above 10")
+	assert.Equal(t, 1, retcode, "ExitOnError and ErrorSeverityLevel = 0, Raiserror above 10")
+	retcode, err = s.runQuery("RAISERROR (N'Testing!' , 5, 1)")
+	assert.NoError(t, err, "ExitOnError and ErrorSeverityLevel = 0, Raiserror below 10")
+	assert.Equal(t, -101, retcode, "ExitOnError and ErrorSeverityLevel = 0, Raiserror below 10")
+}
+
+func TestSqlCmdExitOnError(t *testing.T) {
+	s, buf := setupSqlCmdWithMemoryOutput(t)
+	s.Connect.ExitOnError = true
+	err := runSqlCmd(t, s, []string{"select 1", "GO", ":setvar", "select 2", "GO"})
+	o := buf.buf.String()
+	assert.EqualError(t, err, "Sqlcmd: Error: Syntax error at line 3 near command ':SETVAR'.", "Run should return an error")
+	assert.Equal(t, "1"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol, o, "Only first select should run")
+	assert.Equal(t, 1, s.Exitcode, "s.ExitCode for a syntax error")
+
+	s, buf = setupSqlCmdWithMemoryOutput(t)
+	s.Connect.ExitOnError = true
+	s.Connect.ErrorSeverityLevel = 15
+	s.vars.Set(SQLCMDERRORLEVEL, "14")
+	err = runSqlCmd(t, s, []string{"raiserror(N'13', 13, 1)", "GO", "raiserror(N'14', 14, 1)", "GO", "raiserror(N'15', 15, 1)", "GO", "SELECT 'nope'", "GO"})
+	o = buf.buf.String()
+	assert.NotContains(t, o, "Level 13", "Level 13 should be filtered from the output")
+	assert.NotContains(t, o, "nope", "Last select should not be run")
+	assert.Contains(t, o, "Level 14", "Level 14 should be in the output")
+	assert.Contains(t, o, "Level 15", "Level 15 should be in the output")
+	assert.Equal(t, 15, s.Exitcode, "s.ExitCode for a syntax error")
+	assert.NoError(t, err, "Run should not return an error for a SQL error")
+}
+
+func TestSqlCmdSetErrorLevel(t *testing.T) {
+	s, _ := setupSqlCmdWithMemoryOutput(t)
+	s.Connect.ErrorSeverityLevel = 15
+	err := runSqlCmd(t, s, []string{"select bad as bad", "GO", "select 1", "GO"})
+	assert.NoError(t, err, "runSqlCmd should have no error")
+	assert.Equal(t, 16, s.Exitcode, "Select error should be the exit code")
+}
+
+func runSqlCmd(t testing.TB, s *Sqlcmd, lines []string) error {
+	t.Helper()
+	i := 0
+	s.batch.read = func() (string, error) {
+		if i < len(lines) {
+			index := i
+			i++
+			return lines[index], nil
+		}
+		return "", io.EOF
+	}
+	return s.Run(false, false)
+}
+
 func setupSqlCmdWithMemoryOutput(t testing.TB) (*Sqlcmd, *memoryBuffer) {
 	v := InitializeVariables(true)
 	v.Set(SQLCMDMAXVARTYPEWIDTH, "0")

pkg/sqlcmd/variables.go

@@ -163,6 +163,11 @@ func (v Variables) RowsBetweenHeaders() int64 {
 	return h
 }
 
+// ErrorLevel controls the minimum level of errors that are printed
+func (v Variables) ErrorLevel() int64 {
+	return mustValue(v[SQLCMDERRORLEVEL])
+}
+
 func mustValue(val string) int64 {
 	var n int64
 	_, err := fmt.Sscanf(val, "%d", &n)