Fixing issues with azuresql:// URLs

dburl.go

@@ -131,6 +131,10 @@ func Parse(urlstr string) (*URL, error) {
 	if u.DSN, err = scheme.Generator(u); err != nil {
 		return nil, err
 	}
+	// fix actual driver
+	if scheme.Actual != nil {
+		u.Driver = scheme.Actual(u)
+	}
 	return u, nil
 }
 

dburl_test.go

@@ -124,8 +124,10 @@ func TestParse(t *testing.T) {
 		{`mssql://user:pass@localhost/dbname`, `sqlserver`, `sqlserver://user:pass@localhost/?database=dbname`, ``},
 		{`mssql://user@localhost/service/dbname`, `sqlserver`, `sqlserver://user@localhost/service?database=dbname`, ``},
 		{`mssql://user:!234%23$@localhost:1580/dbname`, `sqlserver`, `sqlserver://user:%21234%23$@localhost:1580/?database=dbname`, ``},
-		{`mssql://user:!234%23$@localhost:1580/service/dbname?fedauth=true`, `sqlserver`, `azuresql://user:%21234%23$@localhost:1580/service?database=dbname&fedauth=true`, ``},
-		{`azuresql://user:pass@localhost:100/dbname`, `sqlserver`, `azuresql://user:pass@localhost:100/?database=dbname`, ``},
+		{`mssql://user:!234%23$@localhost:1580/service/dbname?fedauth=true`, `azuresql`, `azuresql://user:%21234%23$@localhost:1580/service?database=dbname&fedauth=true`, ``},
+		{`azuresql://user:pass@localhost:100/dbname`, `azuresql`, `azuresql://user:pass@localhost:100/?database=dbname`, ``},
+		{`sqlserver://xxx.database.windows.net?database=xxx&fedauth=ActiveDirectoryMSI`, `azuresql`, `azuresql://xxx.database.windows.net?database=xxx&fedauth=ActiveDirectoryMSI`, ``},
+		{`azuresql://xxx.database.windows.net/dbname?fedauth=ActiveDirectoryMSI`, `azuresql`, `azuresql://xxx.database.windows.net/?database=dbname&fedauth=ActiveDirectoryMSI`, ``},
 		{
 			`adodb://Microsoft.ACE.OLEDB.12.0?Extended+Properties=%22Text%3BHDR%3DNO%3BFMT%3DDelimited%22`, `adodb`, // 30
 			`Data Source=.;Extended Properties="Text;HDR=NO;FMT=Delimited";Provider=Microsoft.ACE.OLEDB.12.0`, ``,
@@ -209,19 +211,17 @@ func TestParse(t *testing.T) {
 	}
 	for i, test := range tests {
 		u, err := Parse(test.s)
-		if err != nil {
-			t.Errorf("test %d expected no error, got: %v", i, err)
-			continue
-		}
-		if u.Driver != test.d {
+		switch {
+		case err != nil:
+			t.Fatalf("test %d expected no error, got: %v", i, err)
+		case u.Driver != test.d:
 			t.Errorf("test %d expected driver %q, got: %q", i, test.d, u.Driver)
-		}
-		if u.DSN != test.exp {
+		case u.DSN != test.exp:
 			_, err := os.Stat(test.path)
 			if test.path != "" && err != nil && os.IsNotExist(err) {
 				t.Logf("test %d expected dsn %q, got: %q -- ignoring because `%s` does not exist", i, test.exp, u.DSN, test.path)
 			} else {
-				t.Errorf("test %d expected dsn %q, got: %q", i, test.exp, u.DSN)
+				t.Errorf("test %d expected:\n%q\ngot:\n%q", i, test.exp, u.DSN)
 			}
 		}
 	}

dsn.go

@@ -23,20 +23,19 @@ var Stat = func(name string) (fs.FileInfo, error) {
 // passed URL.
 func GenScheme(scheme string) func(*URL) (string, error) {
 	return func(u *URL) (string, error) {
-		host := u.Host
-		if host == "" {
-			host = "localhost"
-		}
 		z := &url.URL{
 			Scheme:   scheme,
 			Opaque:   u.Opaque,
 			User:     u.User,
-			Host:     host,
+			Host:     u.Host,
 			Path:     u.Path,
 			RawPath:  u.RawPath,
 			RawQuery: u.RawQuery,
 			Fragment: u.Fragment,
 		}
+		if z.Host == "" {
+			z.Host = "localhost"
+		}
 		return z.String(), nil
 	}
 }
@@ -539,21 +538,17 @@ func GenSpanner(u *URL) (string, error) {
 
 // GenSqlserver generates a sqlserver DSN from the passed URL.
 func GenSqlserver(u *URL) (string, error) {
-	host := u.Host
-	if host == "" {
-		host = "localhost"
-	}
 	z := &url.URL{
-		Scheme:   "sqlserver",
+		Scheme:   sqlserverDriver(u),
 		Opaque:   u.Opaque,
 		User:     u.User,
-		Host:     host,
+		Host:     u.Host,
 		Path:     u.Path,
 		RawQuery: u.RawQuery,
 		Fragment: u.Fragment,
 	}
-	if strings.ToLower(u.Scheme) == "azuresql" || z.Query().Has("fedauth") {
-		z.Scheme = "azuresql"
+	if z.Host == "" {
+		z.Host = "localhost"
 	}
 	v := strings.Split(strings.TrimPrefix(z.Path, "/"), "/")
 	if n, q := len(v), z.Query(); !q.Has("database") && n != 0 && len(v[0]) != 0 {
@@ -563,6 +558,16 @@ func GenSqlserver(u *URL) (string, error) {
 	return z.String(), nil
 }
 
+// sqlserverDriver returns the driver used for a Microsoft SQL Server URL.
+func sqlserverDriver(u *URL) string {
+	switch {
+	case u.Query().Has("fedauth"),
+		strings.Contains(strings.ToLower(u.OriginalScheme), "azuresql"):
+		return "azuresql"
+	}
+	return "sqlserver"
+}
+
 // GenTableStore generates a tablestore DSN from the passed URL.
 func GenTableStore(u *URL) (string, error) {
 	var transport string