Add --database to sqlcmd config cs (#280)

Author: stuartpa

cmd/modern/main.go

@@ -19,6 +19,7 @@ import (
 	"github.com/microsoft/go-sqlcmd/internal/io/file"
 	"github.com/microsoft/go-sqlcmd/internal/output"
 	"github.com/microsoft/go-sqlcmd/internal/output/verbosity"
+	"github.com/microsoft/go-sqlcmd/internal/pal"
 	"github.com/microsoft/go-sqlcmd/pkg/sqlcmd"
 	"github.com/spf13/cobra"
 	"path"
@@ -142,7 +143,7 @@ func initializeCallback() {
 // To aid debugging issues, if the logging level is > 2 (e.g. --verbosity 3 or --verbosity 4), we
 // panic which outputs a stacktrace.
 func checkErr(err error) {
-	if rootCmd.loggingLevel > 2 {
+	if rootCmd != nil && rootCmd.loggingLevel > 2 {
 		if err != nil {
 			panic(err)
 		}
@@ -155,7 +156,7 @@ func checkErr(err error) {
 // to make progress.  displayHints is injected into dependencies (helpers etc.)
 func displayHints(hints []string) {
 	if len(hints) > 0 {
-		outputter.Infof("%vHINT:", sqlcmd.SqlcmdEol)
+		outputter.Infof("%vHINT:", pal.LineBreak())
 		for i, hint := range hints {
 			outputter.Infof("  %d. %v", i+1, hint)
 		}

cmd/modern/main_test.go

@@ -5,11 +5,11 @@ package main
 
 import (
 	"errors"
+	"github.com/microsoft/go-sqlcmd/internal/buffer"
 	"github.com/microsoft/go-sqlcmd/internal/cmdparser"
 	"github.com/microsoft/go-sqlcmd/internal/cmdparser/dependency"
 	"github.com/microsoft/go-sqlcmd/internal/output"
 	"github.com/microsoft/go-sqlcmd/internal/pal"
-	"github.com/microsoft/go-sqlcmd/internal/test"
 	"github.com/stretchr/testify/assert"
 	"os"
 	"testing"
@@ -26,14 +26,15 @@ func TestInitializeCallback(t *testing.T) {
 }
 
 func TestDisplayHints(t *testing.T) {
-	buf := test.NewMemoryBuffer()
-	defer checkErr(buf.Close())
+	buf := buffer.NewMemoryBuffer()
 	outputter = output.New(output.Options{StandardWriter: buf})
 	displayHints([]string{"This is a hint"})
 	assert.Equal(t, pal.LineBreak()+
 		"HINT:"+
 		pal.LineBreak()+
 		"  1. This is a hint"+pal.LineBreak()+pal.LineBreak(), buf.String())
+	err := buf.Close()
+	checkErr(err)
 }
 
 func TestCheckErr(t *testing.T) {

cmd/modern/root/config/connection-strings.go

@@ -6,6 +6,8 @@ package config
 import (
 	"fmt"
 	"github.com/microsoft/go-sqlcmd/cmd/modern/sqlconfig"
+	"github.com/microsoft/go-sqlcmd/internal/container"
+	"github.com/microsoft/go-sqlcmd/internal/sql"
 	"strings"
 
 	"github.com/microsoft/go-sqlcmd/internal/cmdparser"
@@ -17,6 +19,8 @@ import (
 // ConnectionStrings implements the `sqlcmd config connection-strings` command
 type ConnectionStrings struct {
 	cmdparser.Cmd
+
+	database string
 }
 
 func (c *ConnectionStrings) DefineCommand(...cmdparser.CommandOptions) {
@@ -36,6 +40,13 @@ func (c *ConnectionStrings) DefineCommand(...cmdparser.CommandOptions) {
 	}
 
 	c.Cmd.DefineCommand(options)
+
+	c.AddFlag(cmdparser.FlagOptions{
+		String:        &c.database,
+		Name:          "database",
+		DefaultString: "",
+		Shorthand:     "d",
+		Usage:         "Database for the connection string (default is taken from the T/SQL login)"})
 }
 
 // run generates connection strings for the current context in multiple formats.
@@ -48,12 +59,27 @@ func (c *ConnectionStrings) run() {
 		"ADO.NET": "Server=tcp:%s,%d;Initial Catalog=%s;Persist Security Info=False;User ID=%s;Password=%s;MultipleActiveResultSets=False;Encrypt=True;TrustServerCertificate=%s;Connection Timeout=30;",
 		"JDBC":    "jdbc:sqlserver://%s:%d;database=%s;user=%s;password=%s;encrypt=true;trustServerCertificate=%s;loginTimeout=30;",
 		"ODBC":    "Driver={ODBC Driver 18 for SQL Server};Server=tcp:%s,%d;Database=%s;Uid=%s;Pwd=%s;Encrypt=yes;TrustServerCertificate=%s;Connection Timeout=30;",
-		"GO":      "sqlserver://%s:%s@%s,%d?database=master;encrypt=true;trustServerCertificate=%s;dial+timeout=30",
+		"GO":      "sqlserver://%s:%s@%s,%d?database=%s;encrypt=true;trustServerCertificate=%s;dial+timeout=30",
 		"SQLCMD":  "sqlcmd -S %s,%d -U %s",
 	}
 
 	endpoint, user := config.CurrentContext()
 
+	if c.database == "" {
+		if endpoint.AssetDetails != nil && endpoint.AssetDetails.ContainerDetails != nil {
+			controller := container.NewController()
+			if controller.ContainerRunning(endpoint.AssetDetails.ContainerDetails.Id) {
+				s := sql.New(sql.SqlOptions{})
+				s.Connect(endpoint, user, sql.ConnectOptions{Interactive: false})
+				c.database = s.ScalarString("PRINT DB_NAME()")
+			} else {
+				c.database = "master"
+			}
+		} else {
+			c.database = "master"
+		}
+	}
+
 	if user != nil {
 		for k, v := range connectionStringFormats {
 			if k == "GO" {
@@ -63,23 +89,25 @@ func (c *ConnectionStrings) run() {
 					secret.Decode(user.BasicAuth.Password, user.BasicAuth.PasswordEncrypted),
 					endpoint.EndpointDetails.Address,
 					endpoint.EndpointDetails.Port,
+					c.database,
 					c.stringForBoolean(c.trustServerCertificate(endpoint), k))
 			} else if k == "SQLCMD" {
 				format := pal.CmdLineWithEnvVars(
 					[]string{"SQLCMDPASSWORD=%s"},
-					"sqlcmd -S %s,%d -U %s",
+					"sqlcmd -S %s,%d -U %s -d %s",
 				)
 
 				connectionStringFormats[k] = fmt.Sprintf(format,
 					secret.Decode(user.BasicAuth.Password, user.BasicAuth.PasswordEncrypted),
 					endpoint.EndpointDetails.Address,
 					endpoint.EndpointDetails.Port,
-					user.BasicAuth.Username)
+					user.BasicAuth.Username,
+					c.database)
 			} else {
 				connectionStringFormats[k] = fmt.Sprintf(v,
 					endpoint.EndpointDetails.Address,
 					endpoint.EndpointDetails.Port,
-					"master",
+					c.database,
 					user.BasicAuth.Username,
 					secret.Decode(user.BasicAuth.Password, user.BasicAuth.PasswordEncrypted),
 					c.stringForBoolean(c.trustServerCertificate(endpoint), k))

cmd/modern/root/config/connection-strings_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/microsoft/go-sqlcmd/internal"
 	"github.com/microsoft/go-sqlcmd/internal/cmdparser"
 	"github.com/microsoft/go-sqlcmd/internal/output"
+	"github.com/stretchr/testify/assert"
 	"os"
 	"testing"
 )
@@ -48,5 +49,10 @@ func TestConnectionStrings(t *testing.T) {
 	cmdparser.TestCmd[*AddUser]("--username user")
 	cmdparser.TestCmd[*AddContext]("--endpoint endpoint2 --user user")
 
-	cmdparser.TestCmd[*ConnectionStrings]()
+	result := cmdparser.TestCmd[*ConnectionStrings]()
+	assert.Contains(t, result, "database=master")
+
+	result = cmdparser.TestCmd[*ConnectionStrings]("--database tempdb")
+	assert.NotContains(t, result, "database=master")
+	assert.Contains(t, result, "database=tempdb")
 }

internal/test/memory-buffer.go renamed to internal/buffer/memory-buffer.go

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-package test
+package buffer
 
 import "bytes"
 
@@ -16,6 +16,8 @@ func (b *MemoryBuffer) Write(p []byte) (n int, err error) {
 }
 
 func (b *MemoryBuffer) Close() error {
+	b.buf = nil
+
 	return nil
 }
 

internal/test/memory-buffer_test.go renamed to internal/buffer/memory-buffer_test.go

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-package test
+package buffer
 
 import (
 	"testing"

internal/cmdparser/test.go

@@ -5,6 +5,7 @@ package cmdparser
 
 import (
 	"github.com/microsoft/go-sqlcmd/internal"
+	"github.com/microsoft/go-sqlcmd/internal/buffer"
 	"github.com/microsoft/go-sqlcmd/internal/cmdparser/dependency"
 	"github.com/microsoft/go-sqlcmd/internal/config"
 	"github.com/microsoft/go-sqlcmd/internal/output"
@@ -44,21 +45,26 @@ func TestSetup(t *testing.T) {
 }
 
 // Run a command expecing it to pass, passing in any supplied args (args are split on " " (space))
-func TestCmd[T PtrAsReceiverWrapper[pointerType], pointerType any](args ...string) {
-	err := testCmd[T](args...)
+func TestCmd[T PtrAsReceiverWrapper[pointerType], pointerType any](args ...string) string {
+	result, err := testCmd[T](args...)
 
-	// DEVNOTE: I don't think the code will ever get here (c.Command().Execute() will
-	// always panic first. This is here to silence code checkers, that require the err return
-	// variable be checked.
 	if err != nil {
+
+		// DEVNOTE: I don't think the code will ever get here (c.Command().Execute() will
+		// always panic first. This is here to silence code checkers, that require the err return
+		// variable be used.
 		panic(err)
 	}
+	return result
 }
 
-func testCmd[T PtrAsReceiverWrapper[pointerType], pointerType any](args ...string) error {
+func testCmd[T PtrAsReceiverWrapper[pointerType], pointerType any](args ...string) (result string, err error) {
+	buf := buffer.NewMemoryBuffer()
+	defer func() { buf.Close() }()
 	c := New[T](dependency.Options{
 		Output: output.New(output.Options{
-			LoggingLevel: verbosity.Trace}),
+			StandardWriter: buf,
+			LoggingLevel:   verbosity.Trace}),
 	})
 	c.DefineCommand()
 	if len(args) > 1 {
@@ -68,8 +74,8 @@ func testCmd[T PtrAsReceiverWrapper[pointerType], pointerType any](args ...strin
 	} else {
 		c.SetArgsForUnitTesting([]string{})
 	}
-	err := c.Command().Execute()
-	return err
+	err = c.Command().Execute()
+	return buf.String(), err
 }
 
 // splitStringIntoArgsSlice uses a regular expression that matches either a

internal/sql/interface.go

@@ -10,6 +10,7 @@ import (
 type Sql interface {
 	Connect(endpoint Endpoint, user *User, options ConnectOptions)
 	Query(text string)
+	ScalarString(query string) string
 }
 
 type ConnectOptions struct {

internal/sql/mock.go

@@ -16,3 +16,7 @@ func (m *mock) Connect(
 // Query is a mock implementation used to speed up unit testing of other units
 func (m *mock) Query(text string) {
 }
+
+func (m *mock) ScalarString(query string) string {
+	return ""
+}

internal/sql/mssql.go

@@ -5,8 +5,10 @@ package sql
 
 import (
 	"fmt"
+	"github.com/microsoft/go-sqlcmd/internal/buffer"
 	"github.com/microsoft/go-sqlcmd/pkg/console"
 	"os"
+	"strings"
 
 	"github.com/microsoft/go-sqlcmd/cmd/modern/sqlconfig"
 	"github.com/microsoft/go-sqlcmd/pkg/sqlcmd"
@@ -77,3 +79,18 @@ func (m *mssql) Query(text string) {
 		checkErr(err)
 	}
 }
+
+func (m *mssql) ScalarString(query string) string {
+	buf := buffer.NewMemoryBuffer()
+	defer func() { _ = buf.Close() }()
+
+	m.sqlcmd.Query = query
+	m.sqlcmd.SetOutput(buf)
+	m.sqlcmd.SetError(os.Stderr)
+
+	trace("Running query: %v", query)
+	err := m.sqlcmd.Run(true, false)
+	checkErr(err)
+
+	return strings.TrimRight(buf.String(), "\r\n")
+}