implement messages based on private mssqldb (#41)

* implement messages based on private mssqldb

* switch back to the primary driver source

Author: shueybubbles

.vscode/launch.json

@@ -26,5 +26,13 @@
             "program": "${workspaceFolder}/cmd/sqlcmd",
             "args" : ["-Q", "EXIT(select 100 as Count)"],
         },
+        {
+            "name" : "Run file query",
+            "type" : "go",
+            "request": "launch",
+            "mode" : "auto",
+            "program": "${workspaceFolder}/cmd/sqlcmd",
+            "args" : ["-i", "testdata\\select100.sql"],
+        },
     ]
 }

cmd/sqlcmd/main_test.go

@@ -13,6 +13,8 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+const oneRowAffected = "(1 row affected)"
+
 func newKong(t *testing.T, cli interface{}, options ...kong.Option) *kong.Kong {
 	t.Helper()
 	options = append([]kong.Option{
@@ -122,7 +124,7 @@ func TestRunInputFiles(t *testing.T) {
 	assert.Equal(t, 0, exitCode, "exitCode")
 	bytes, err := os.ReadFile(o.Name())
 	if assert.NoError(t, err, "os.ReadFile") {
-		assert.Equal(t, "100"+sqlcmd.SqlcmdEol+sqlcmd.SqlcmdEol+"100"+sqlcmd.SqlcmdEol+sqlcmd.SqlcmdEol, string(bytes), "Incorrect output from run")
+		assert.Equal(t, "100"+sqlcmd.SqlcmdEol+sqlcmd.SqlcmdEol+oneRowAffected+sqlcmd.SqlcmdEol+"100"+sqlcmd.SqlcmdEol+sqlcmd.SqlcmdEol+oneRowAffected+sqlcmd.SqlcmdEol, string(bytes), "Incorrect output from run")
 	}
 }
 
@@ -148,7 +150,7 @@ func TestQueryAndExit(t *testing.T) {
 	assert.Equal(t, 0, exitCode, "exitCode")
 	bytes, err := os.ReadFile(o.Name())
 	if assert.NoError(t, err, "os.ReadFile") {
-		assert.Equal(t, "100 val2"+sqlcmd.SqlcmdEol+sqlcmd.SqlcmdEol, string(bytes), "Incorrect output from run")
+		assert.Equal(t, "100 val2"+sqlcmd.SqlcmdEol+sqlcmd.SqlcmdEol+oneRowAffected+sqlcmd.SqlcmdEol, string(bytes), "Incorrect output from run")
 	}
 }
 
@@ -174,7 +176,7 @@ func TestAzureAuth(t *testing.T) {
 	assert.Equal(t, 0, exitCode, "exitCode")
 	bytes, err := os.ReadFile(o.Name())
 	if assert.NoError(t, err, "os.ReadFile") {
-		assert.Equal(t, "AZURE"+sqlcmd.SqlcmdEol+sqlcmd.SqlcmdEol, string(bytes), "Incorrect output from run")
+		assert.Equal(t, "AZURE"+sqlcmd.SqlcmdEol+sqlcmd.SqlcmdEol+oneRowAffected+sqlcmd.SqlcmdEol, string(bytes), "Incorrect output from run")
 	}
 }
 

cmd/sqlcmd/testdata/create100db.sql

@@ -0,0 +1,30 @@
+create database db100
+go
+
+create database db101
+go
+
+create database db102
+go
+create database db103
+go
+create database db104
+go
+create database db105
+go
+create database db106
+go
+create database db107
+go
+create database db108
+go
+create database db109
+go
+create database db110
+go
+create database db111
+go
+create database db112
+go
+create database db113
+go

cmd/sqlcmd/testdata/drop100db.txt

@@ -0,0 +1,30 @@
+drop database db100
+go
+
+drop database db101
+go
+
+drop database db102
+go
+drop database db103
+go
+drop database db104
+go
+drop database db105
+go
+drop database db106
+go
+drop database db107
+go
+drop database db108
+go
+drop database db109
+go
+drop database db110
+go
+drop database db111
+go
+drop database db112
+go
+drop database db113
+go

go.mod

@@ -8,8 +8,10 @@ require (
 	github.com/alecthomas/kong v0.2.18-0.20210621093454-54558f65e86f
 	github.com/chzyer/logex v1.1.10 // indirect
 	github.com/chzyer/test v0.0.0-20210722231415-061457976a23 // indirect
-	github.com/denisenkom/go-mssqldb v0.10.0
+	github.com/denisenkom/go-mssqldb v0.12.0
 	github.com/gohxs/readline v0.0.0-20171011095936-a780388e6e7c
+	github.com/golang-sql/sqlexp v0.0.0-20170517235910-f1bb20e5a188
 	github.com/google/uuid v1.2.0
 	github.com/stretchr/testify v1.7.0
 )
+

go.sum

@@ -4,8 +4,6 @@ github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.11.0 h1:OYa9vmRX2XC5GXRAzegg
 github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.11.0/go.mod h1:HcM1YX14R7CJcghJGOYCgdezslRSVzqwLf/q+4Y2r/0=
 github.com/Azure/azure-sdk-for-go/sdk/internal v0.7.0 h1:v9p9TfTbf7AwNb5NYQt7hI41IfPoLFiFkLtb+bmGjT0=
 github.com/Azure/azure-sdk-for-go/sdk/internal v0.7.0/go.mod h1:yqy467j36fJxcRV2TzfVZ1pCb5vxm4BtZPUdYWe/Xo8=
-github.com/alecthomas/kong v0.2.17 h1:URDISCI96MIgcIlQyoCAlhOmrSw6pZScBNkctg8r0W0=
-github.com/alecthomas/kong v0.2.17/go.mod h1:ka3VZ8GZNPXv9Ov+j4YNLkI8mTuhXyr/0ktSlqIydQQ=
 github.com/alecthomas/kong v0.2.18-0.20210621093454-54558f65e86f h1:VgRM6/wqZIB1D9W3XMllm/wplTmPgI5yvCHUXEsmKps=
 github.com/alecthomas/kong v0.2.18-0.20210621093454-54558f65e86f/go.mod h1:ka3VZ8GZNPXv9Ov+j4YNLkI8mTuhXyr/0ktSlqIydQQ=
 github.com/chzyer/logex v1.1.10 h1:Swpa1K6QvQznwJRcfTfQJmTE72DqScAa40E+fbHEXEE=
@@ -15,13 +13,15 @@ github.com/chzyer/test v0.0.0-20210722231415-061457976a23/go.mod h1:Q3SI9o4m/ZMn
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/denisenkom/go-mssqldb v0.10.0 h1:QykgLZBorFE95+gO3u9esLd0BmbvpWp0/waNNZfHBM8=
-github.com/denisenkom/go-mssqldb v0.10.0/go.mod h1:xbL0rPBG9cCiLr28tMa8zpbdarY27NDyej4t/EjAShU=
+github.com/denisenkom/go-mssqldb v0.12.0 h1:VtrkII767ttSPNRfFekePK3sctr+joXgO58stqQbtUA=
+github.com/denisenkom/go-mssqldb v0.12.0/go.mod h1:iiK0YP1ZeepvmBQk/QpLEhhTNJgfzrpArPY/aFvc9yU=
 github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ=
 github.com/gohxs/readline v0.0.0-20171011095936-a780388e6e7c h1:yE35fKFwcelIte3q5q1/cPiY7pI7vvf5/j/0ddxNCKs=
 github.com/gohxs/readline v0.0.0-20171011095936-a780388e6e7c/go.mod h1:9S/fKAutQ6wVHqm1jnp9D9sc5hu689s9AaTWFS92LaU=
 github.com/golang-sql/civil v0.0.0-20190719163853-cb61b32ac6fe h1:lXe2qZdvpiX5WZkZR4hgp4KJVfY3nMkvmwbVkpv1rVY=
 github.com/golang-sql/civil v0.0.0-20190719163853-cb61b32ac6fe/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0=
+github.com/golang-sql/sqlexp v0.0.0-20170517235910-f1bb20e5a188 h1:+eHOFJl1BaXrQxKX+T06f78590z4qA2ZzBTqahsKSE4=
+github.com/golang-sql/sqlexp v0.0.0-20170517235910-f1bb20e5a188/go.mod h1:vXjM/+wXQnTPR4KqTKDgJukSZ6amVRtWMPEjE6sQoK8=
 github.com/google/uuid v1.2.0 h1:qJYtXnJRWmpe7m/3XlyhrsLrEURqHRM2kxzoxXqyUDs=
 github.com/google/uuid v1.2.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/modocache/gover v0.0.0-20171022184752-b58185e213c5/go.mod h1:caMODM3PzxT8aQXRPkAt8xlV/e7d7w8GM5g0fa5F0D8=
@@ -35,7 +35,6 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
-golang.org/x/crypto v0.0.0-20190325154230-a5d413f7728c/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897 h1:pLI5jrR7OSLijeIDcmRxNmw2api+jEfxLoykJVice/E=
 golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=

pkg/sqlcmd/format.go

@@ -166,7 +166,9 @@ func (f *sqlCmdFormatterType) AddRow(row *sql.Rows) string {
 }
 
 // Writes a non-error message to the designated message writer
-func (f *sqlCmdFormatterType) AddMessage(string) {}
+func (f *sqlCmdFormatterType) AddMessage(msg string) {
+	f.mustWriteOut(msg + SqlcmdEol)
+}
 
 // Writes an error to the designated err Writer
 func (f *sqlCmdFormatterType) AddError(err error) {
@@ -273,7 +275,9 @@ func calcColumnDetails(cols []*sql.ColumnType, fixed int64, variable int64) (col
 		} else {
 			columnDetails[i].displayWidth = length
 		}
-		switch c.DatabaseTypeName() {
+		typeName := c.DatabaseTypeName()
+
+		switch typeName {
 		// Types with 0 size from sql.ColumnType
 		case "BIT":
 			columnDetails[i].leftJustify = false

pkg/sqlcmd/sqlcmd.go

@@ -5,6 +5,7 @@ package sqlcmd
 
 import (
 	"bufio"
+	"context"
 	"database/sql"
 	"database/sql/driver"
 	"errors"
@@ -20,6 +21,7 @@ import (
 
 	mssql "github.com/denisenkom/go-mssqldb"
 	"github.com/gohxs/readline"
+	"github.com/golang-sql/sqlexp"
 )
 
 var (
@@ -450,33 +452,61 @@ func (s *Sqlcmd) sqlAuthentication() bool {
 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)
+	ctx := context.Background()
+	retmsg := &sqlexp.ReturnMessage{}
+	rows, qe := s.db.QueryContext(ctx, query, retmsg)
 	if qe != nil {
 		s.Format.AddError(qe)
 	}
 	var err error
 	var cols []*sql.ColumnType
 	results := true
+	first := 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 {
+		msg := retmsg.Message(ctx)
+		switch m := msg.(type) {
+		case sqlexp.MsgNotice:
+			s.Format.AddMessage(m.Message)
+		case sqlexp.MsgError:
+			s.Format.AddError(m.Error)
+		case sqlexp.MsgRowsAffected:
+			if m.Count == 1 {
+				s.Format.AddMessage("(1 row affected)")
+			} else {
+				s.Format.AddMessage(fmt.Sprintf("(%d rows affected)", m.Count))
+			}
+		case sqlexp.MsgNextResultSet:
+			results = rows.NextResultSet()
+			if err = rows.Err(); err != nil {
+				retcode = -100
+				s.Format.AddError(err)
+			}
+			if results {
+				first = true
+			}
+		case sqlexp.MsgNext:
+			inresult := rows.Next()
+			for inresult {
+				if first {
+					first = false
+					cols, err = rows.ColumnTypes()
+					if err != nil {
+						retcode = -100
+						s.Format.AddError(err)
+					} else {
+						s.Format.BeginResultSet(cols)
+					}
+				}
 				col1 := s.Format.AddRow(rows)
-				active = rows.Next()
-				if !active {
+				inresult = rows.Next()
+				if !inresult {
 					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
@@ -485,11 +515,6 @@ func (s *Sqlcmd) runQuery(query string) int {
 			}
 			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

@@ -16,6 +16,8 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+const oneRowAffected = "(1 row affected)"
+
 func TestConnectionStringFromSqlCmd(t *testing.T) {
 	type connectionStringTest struct {
 		settings         *ConnectSettings
@@ -159,7 +161,7 @@ func TestIncludeFileNoExecutions(t *testing.T) {
 			assert.Equal(t, "select 'string' as title", s.batch.String(), "s.batch.String() after IncludeFile twobatchnoendinggo.sql false")
 			bytes, err := os.ReadFile(file.Name())
 			if assert.NoError(t, err, "os.ReadFile") {
-				assert.Equal(t, "100"+SqlcmdEol+SqlcmdEol+"string"+SqlcmdEol+SqlcmdEol+"100"+SqlcmdEol+SqlcmdEol, string(bytes), "Incorrect output from Run")
+				assert.Equal(t, "100"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol+"string"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol+"100"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol, string(bytes), "Incorrect output from Run")
 			}
 		}
 	}
@@ -177,7 +179,7 @@ func TestIncludeFileProcessAll(t *testing.T) {
 		assert.Equal(t, "", s.batch.String(), "s.batch.String() after IncludeFile twobatchwithgo.sql true")
 		bytes, err := os.ReadFile(file.Name())
 		if assert.NoError(t, err, "os.ReadFile") {
-			assert.Equal(t, "100"+SqlcmdEol+SqlcmdEol+"string"+SqlcmdEol+SqlcmdEol, string(bytes), "Incorrect output from Run")
+			assert.Equal(t, "100"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol+"string"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol, string(bytes), "Incorrect output from Run")
 		}
 		file, err = os.CreateTemp("", "sqlcmdout")
 		defer os.Remove(file.Name())
@@ -189,7 +191,7 @@ func TestIncludeFileProcessAll(t *testing.T) {
 			assert.Equal(t, "", s.batch.String(), "s.batch.String() after IncludeFile twobatchnoendinggo.sql true")
 			bytes, err := os.ReadFile(file.Name())
 			if assert.NoError(t, err, "os.ReadFile") {
-				assert.Equal(t, "100"+SqlcmdEol+SqlcmdEol+"string"+SqlcmdEol+SqlcmdEol, string(bytes), "Incorrect output from Run")
+				assert.Equal(t, "100"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol+"string"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol, string(bytes), "Incorrect output from Run")
 			}
 		}
 	}
@@ -232,7 +234,7 @@ func TestExitInitialQuery(t *testing.T) {
 	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 2100"+SqlcmdEol+SqlcmdEol+oneRowAffected+SqlcmdEol, o, "Output")
 		assert.Equal(t, 1200, s.Exitcode, "ExitCode")
 	}