From 82b9a8c0c1422d2ef327b88a3c98523a995ac885 Mon Sep 17 00:00:00 2001
From: Kenneth Shaw <ken@usql.app>
Date: Wed, 16 Aug 2023 05:34:04 +0700
Subject: [PATCH] Fixing issues with azuresql:// URLs

---
 dburl.go      |   4 +
 dburl_test.go |  20 +--
 dsn.go        |  31 +++--
 scheme.go     | 361 +++++++++++++++++++++++++++++++++++++++++++-------
 4 files changed, 347 insertions(+), 69 deletions(-)

diff --git a/dburl.go b/dburl.go
index 69962b2..75cc708 100644
--- a/dburl.go
+++ b/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
 }
 
diff --git a/dburl_test.go b/dburl_test.go
index d907e7f..7b121c2 100644
--- a/dburl_test.go
+++ b/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)
 			}
 		}
 	}
diff --git a/dsn.go b/dsn.go
index a3eaeeb..e9a9ce4 100644
--- a/dsn.go
+++ b/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
diff --git a/scheme.go b/scheme.go
index 0523124..8cd8df9 100644
--- a/scheme.go
+++ b/scheme.go
@@ -43,61 +43,329 @@ type Scheme struct {
 	//
 	// Used for "wire compatible" driver schemes.
 	Override string
+	// Actual is a func that can be used to inspect the URL that changes the
+	// actual Driver name returned. Used for drivers that have registered
+	// multiple driver names using a single package import (ie, Microsoft's SQL
+	// Server driver).
+	Actual func(*URL) string
 }
 
 // BaseSchemes returns the supported base schemes.
 func BaseSchemes() []Scheme {
 	return []Scheme{
 		// core databases
-		{"mysql", GenMysql, TransportTCP | TransportUDP | TransportUnix, false, []string{"mariadb", "maria", "percona", "aurora"}, ""},
-		{"oracle", GenFromURL("oracle://localhost:1521"), 0, false, []string{"ora", "oci", "oci8", "odpi", "odpi-c"}, ""},
-		{"postgres", GenPostgres, TransportUnix, false, []string{"pg", "postgresql", "pgsql"}, ""},
-		{"sqlite3", GenOpaque, 0, true, []string{"sqlite", "file"}, ""},
-		{"sqlserver", GenSqlserver, 0, false, []string{"ms", "mssql", "azuresql"}, ""},
+		{
+			"mysql",
+			GenMysql, TransportTCP | TransportUDP | TransportUnix,
+			false,
+			[]string{"mariadb", "maria", "percona", "aurora"},
+			"",
+			nil,
+		},
+		{
+			"oracle",
+			GenFromURL("oracle://localhost:1521"), 0, false,
+			[]string{"ora", "oci", "oci8", "odpi", "odpi-c"},
+			"",
+			nil,
+		},
+		{
+			"postgres",
+			GenPostgres, TransportUnix, false,
+			[]string{"pg", "postgresql", "pgsql"},
+			"",
+			nil,
+		},
+		{
+			"sqlite3",
+			GenOpaque, 0, true,
+			[]string{"sqlite", "file"},
+			"",
+			nil,
+		},
+		{
+			"sqlserver",
+			GenSqlserver, 0, false,
+			[]string{"ms", "mssql", "azuresql"},
+			"",
+			sqlserverDriver,
+		},
 		// wire compatibles
-		{"cockroachdb", GenFromURL("postgres://localhost:26257/?sslmode=disable"), 0, false, []string{"cr", "cockroach", "crdb", "cdb"}, "postgres"},
-		{"memsql", GenMysql, 0, false, nil, "mysql"},
-		{"redshift", GenFromURL("postgres://localhost:5439/"), 0, false, []string{"rs"}, "postgres"},
-		{"tidb", GenMysql, 0, false, nil, "mysql"},
-		{"vitess", GenMysql, 0, false, []string{"vt"}, "mysql"},
+		{
+			"cockroachdb",
+			GenFromURL("postgres://localhost:26257/?sslmode=disable"), 0, false,
+			[]string{"cr", "cockroach", "crdb", "cdb"},
+			"postgres",
+			nil,
+		},
+		{
+			"memsql",
+			GenMysql, 0, false, nil, "mysql",
+			nil,
+		},
+		{
+			"redshift",
+			GenFromURL("postgres://localhost:5439/"), 0, false,
+			[]string{"rs"},
+			"postgres",
+			nil,
+		},
+		{
+			"tidb",
+			GenMysql, 0, false, nil, "mysql",
+			nil,
+		},
+		{
+			"vitess",
+			GenMysql, 0, false,
+			[]string{"vt"},
+			"mysql",
+			nil,
+		},
 		// alternate implementations
-		{"godror", GenGodror, 0, false, []string{"gr"}, ""},
-		{"moderncsqlite", GenOpaque, 0, true, []string{"mq", "modernsqlite"}, ""},
-		{"mymysql", GenMymysql, TransportTCP | TransportUDP | TransportUnix, false, []string{"zm", "mymy"}, ""},
-		{"pgx", GenFromURL("postgres://localhost:5432/"), TransportUnix, false, []string{"px"}, ""},
+		{
+			"godror",
+			GenGodror, 0, false,
+			[]string{"gr"},
+			"",
+			nil,
+		},
+		{
+			"moderncsqlite",
+			GenOpaque, 0, true,
+			[]string{"mq", "modernsqlite"},
+			"",
+			nil,
+		},
+		{
+			"mymysql",
+			GenMymysql, TransportTCP | TransportUDP | TransportUnix, false,
+			[]string{"zm", "mymy"},
+			"",
+			nil,
+		},
+		{
+			"pgx",
+			GenFromURL("postgres://localhost:5432/"), TransportUnix, false,
+			[]string{"px"},
+			"",
+			nil,
+		},
 		// other databases
-		{"adodb", GenAdodb, 0, false, []string{"ado"}, ""},
-		{"awsathena", GenScheme("s3"), 0, false, []string{"s3", "aws", "athena"}, ""},
-		{"avatica", GenFromURL("http://localhost:8765/"), 0, false, []string{"phoenix"}, ""},
-		{"bigquery", GenScheme("bigquery"), 0, false, []string{"bq"}, ""},
-		{"clickhouse", GenFromURL("clickhouse://localhost:9000/"), 0, false, []string{"ch"}, ""},
-		{"cosmos", GenCosmos, 0, false, []string{"cm"}, ""},
-		{"cql", GenCassandra, 0, false, []string{"ca", "cassandra", "datastax", "scy", "scylla"}, ""},
-		{"csvq", GenOpaque, 0, true, []string{"csv", "tsv", "json"}, ""},
-		{"databend", GenDatabend, 0, false, []string{"dd", "bend"}, ""},
-		{"exasol", GenExasol, 0, false, []string{"ex", "exa"}, ""},
-		{"firebirdsql", GenFirebird, 0, false, []string{"fb", "firebird"}, ""},
-		{"flightsql", GenScheme("flightsql"), 0, false, []string{"fl", "flight"}, ""},
-		{"genji", GenOpaque, 0, true, []string{"gj"}, ""},
-		{"h2", GenFromURL("h2://localhost:9092/"), 0, false, nil, ""},
-		{"hdb", GenScheme("hdb"), 0, false, []string{"sa", "saphana", "sap", "hana"}, ""},
-		{"hive", GenSchemeTruncate, 0, false, nil, ""},
-		{"ignite", GenIgnite, 0, false, []string{"ig", "gridgain"}, ""},
-		{"impala", GenScheme("impala"), 0, false, nil, ""},
-		{"maxcompute", GenSchemeTruncate, 0, false, []string{"mc"}, ""},
-		{"n1ql", GenFromURL("http://localhost:9000/"), 0, false, []string{"couchbase"}, ""},
-		{"nzgo", GenPostgres, TransportUnix, false, []string{"nz", "netezza"}, ""},
-		{"odbc", GenOdbc, TransportAny, false, nil, ""},
-		{"oleodbc", GenOleodbc, TransportAny, false, []string{"oo", "ole"}, "adodb"},
-		{"ots", GenTableStore, TransportAny, false, []string{"tablestore"}, ""},
-		{"presto", GenPresto, 0, false, []string{"prestodb", "prestos", "prs", "prestodbs"}, ""},
-		{"ql", GenOpaque, 0, true, []string{"ql", "cznic", "cznicql"}, ""},
-		{"snowflake", GenSnowflake, 0, false, []string{"sf"}, ""},
-		{"spanner", GenSpanner, 0, false, []string{"sp"}, ""},
-		{"tds", GenFromURL("http://localhost:5000/"), 0, false, []string{"ax", "ase", "sapase"}, ""},
-		{"trino", GenPresto, 0, false, []string{"trino", "trinos", "trs"}, ""},
-		{"vertica", GenFromURL("vertica://localhost:5433/"), 0, false, nil, ""},
-		{"voltdb", GenVoltdb, 0, false, []string{"volt", "vdb"}, ""},
+		{
+			"adodb",
+			GenAdodb, 0, false,
+			[]string{"ado"},
+			"",
+			nil,
+		},
+		{
+			"awsathena",
+			GenScheme("s3"), 0, false,
+			[]string{"s3", "aws", "athena"},
+			"",
+			nil,
+		},
+		{
+			"avatica",
+			GenFromURL("http://localhost:8765/"), 0, false,
+			[]string{"phoenix"},
+			"",
+			nil,
+		},
+		{
+			"bigquery",
+			GenScheme("bigquery"), 0, false,
+			[]string{"bq"},
+			"",
+			nil,
+		},
+		{
+			"clickhouse",
+			GenFromURL("clickhouse://localhost:9000/"), 0, false,
+			[]string{"ch"},
+			"",
+			nil,
+		},
+		{
+			"cosmos",
+			GenCosmos, 0, false,
+			[]string{"cm"},
+			"",
+			nil,
+		},
+		{
+			"cql",
+			GenCassandra, 0, false,
+			[]string{"ca", "cassandra", "datastax", "scy", "scylla"},
+			"",
+			nil,
+		},
+		{
+			"csvq",
+			GenOpaque, 0, true,
+			[]string{"csv", "tsv", "json"},
+			"",
+			nil,
+		},
+		{
+			"databend",
+			GenDatabend, 0, false,
+			[]string{"dd", "bend"},
+			"",
+			nil,
+		},
+		{
+			"exasol",
+			GenExasol, 0, false,
+			[]string{"ex", "exa"},
+			"",
+			nil,
+		},
+		{
+			"firebirdsql",
+			GenFirebird, 0, false,
+			[]string{"fb", "firebird"},
+			"",
+			nil,
+		},
+		{
+			"flightsql",
+			GenScheme("flightsql"), 0, false,
+			[]string{"fl", "flight"},
+			"",
+			nil,
+		},
+		{
+			"genji",
+			GenOpaque, 0, true,
+			[]string{"gj"},
+			"",
+			nil,
+		},
+		{
+			"h2",
+			GenFromURL("h2://localhost:9092/"), 0, false, nil, "",
+			nil,
+		},
+		{
+			"hdb",
+			GenScheme("hdb"), 0, false,
+			[]string{"sa", "saphana", "sap", "hana"},
+			"",
+			nil,
+		},
+		{
+			"hive",
+			GenSchemeTruncate, 0, false, nil, "",
+			nil,
+		},
+		{
+			"ignite",
+			GenIgnite, 0, false,
+			[]string{"ig", "gridgain"},
+			"",
+			nil,
+		},
+		{
+			"impala",
+			GenScheme("impala"), 0, false, nil, "",
+			nil,
+		},
+		{
+			"maxcompute",
+			GenSchemeTruncate, 0, false,
+			[]string{"mc"},
+			"",
+			nil,
+		},
+		{
+			"n1ql",
+			GenFromURL("http://localhost:9000/"), 0, false,
+			[]string{"couchbase"},
+			"",
+			nil,
+		},
+		{
+			"nzgo",
+			GenPostgres, TransportUnix, false,
+			[]string{"nz", "netezza"},
+			"",
+			nil,
+		},
+		{
+			"odbc",
+			GenOdbc, TransportAny, false, nil, "",
+			nil,
+		},
+		{
+			"oleodbc",
+			GenOleodbc, TransportAny, false,
+			[]string{"oo", "ole"},
+			"adodb",
+			nil,
+		},
+		{
+			"ots",
+			GenTableStore, TransportAny, false,
+			[]string{"tablestore"},
+			"",
+			nil,
+		},
+		{
+			"presto",
+			GenPresto, 0, false,
+			[]string{"prestodb", "prestos", "prs", "prestodbs"},
+			"",
+			nil,
+		},
+		{
+			"ql",
+			GenOpaque, 0, true,
+			[]string{"ql", "cznic", "cznicql"},
+			"",
+			nil,
+		},
+		{
+			"snowflake",
+			GenSnowflake, 0, false,
+			[]string{"sf"},
+			"",
+			nil,
+		},
+		{
+			"spanner",
+			GenSpanner, 0, false,
+			[]string{"sp"},
+			"",
+			nil,
+		},
+		{
+			"tds",
+			GenFromURL("http://localhost:5000/"), 0, false,
+			[]string{"ax", "ase", "sapase"},
+			"",
+			nil,
+		},
+		{
+			"trino",
+			GenPresto, 0, false,
+			[]string{"trino", "trinos", "trs"},
+			"",
+			nil,
+		},
+		{
+			"vertica",
+			GenFromURL("vertica://localhost:5433/"), 0, false, nil, "",
+			nil,
+		},
+		{
+			"voltdb",
+			GenVoltdb, 0, false,
+			[]string{"volt", "vdb"},
+			"",
+			nil,
+		},
 	}
 }
 
@@ -158,6 +426,7 @@ func Register(scheme Scheme) {
 		Transport: scheme.Transport,
 		Opaque:    scheme.Opaque,
 		Override:  scheme.Override,
+		Actual:    scheme.Actual,
 	}
 	schemeMap[scheme.Driver] = sz
 	// add aliases