Repository: ariga/atlas
Pull request #748 (2022-04-28 11:13:23)

sql/sqlclient: move url parsing logic to drivers


Patch:
diff --git a/cmd/action/apply.go b/cmd/action/apply.go
index d98a76c4..3bcc5301 100644
--- a/cmd/action/apply.go
+++ b/cmd/action/apply.go
@@ -82,14 +82,13 @@ func CmdApplyRun(cmd *cobra.Command, _ []string) {
 	c, err := sqlclient.Open(cmd.Context(), ApplyFlags.URL)
 	cobra.CheckErr(err)
 	defer c.Close()
-	applyRun(cmd.Context(), c, ApplyFlags.URL, ApplyFlags.File, ApplyFlags.DryRun, ApplyFlags.AutoApprove)
+	applyRun(cmd.Context(), c, ApplyFlags.File, ApplyFlags.DryRun, ApplyFlags.AutoApprove)
 }
 
-func applyRun(ctx context.Context, client *sqlclient.Client, url string, file string, dryRun bool, autoApprove bool) {
+func applyRun(ctx context.Context, client *sqlclient.Client, file string, dryRun, autoApprove bool) {
 	schemas := ApplyFlags.Schema
-	if n, err := SchemaNameFromURL(ctx, url); n != "" {
-		cobra.CheckErr(err)
-		schemas = append(schemas, n)
+	if client.URL.Schema != "" {
+		schemas = append(schemas, client.URL.Schema)
 	}
 	realm, err := client.InspectRealm(ctx, &schema.InspectRealmOption{
 		Schemas: schemas,
diff --git a/cmd/action/diff.go b/cmd/action/diff.go
index 20d930d9..bb73d12f 100644
--- a/cmd/action/diff.go
+++ b/cmd/action/diff.go
@@ -52,10 +52,8 @@ func cmdDiffRun(cmd *cobra.Command, flags *diffCmdOpts) {
 	toC, err := sqlclient.Open(cmd.Context(), flags.toURL)
 	cobra.CheckErr(err)
 	defer toC.Close()
-	fromName, err := SchemaNameFromURL(ctx, flags.fromURL)
-	cobra.CheckErr(err)
-	toName, err := SchemaNameFromURL(ctx, flags.toURL)
-	cobra.CheckErr(err)
+	fromName := fromC.URL.Schema
+	toName := toC.URL.Schema
 	fromSchema, err := fromC.InspectSchema(ctx, fromName, nil)
 	cobra.CheckErr(err)
 	toSchema, err := toC.InspectSchema(ctx, toName, nil)
diff --git a/cmd/action/inspect.go b/cmd/action/inspect.go
index dabd5279..e0873bd0 100644
--- a/cmd/action/inspect.go
+++ b/cmd/action/inspect.go
@@ -5,8 +5,6 @@
 package action
 
 import (
-	"context"
-
 	"ariga.io/atlas/sql/schema"
 	"ariga.io/atlas/sql/sqlclient"
 
@@ -62,20 +60,14 @@ func CmdInspectRun(cmd *cobra.Command, _ []string) {
 		schemaCmd.PrintErrln("The Atlas UI is not available in this release.")
 		return
 	}
-	d, err := sqlclient.Open(cmd.Context(), InspectFlags.URL)
+	client, err := sqlclient.Open(cmd.Context(), InspectFlags.URL)
 	cobra.CheckErr(err)
-	defer d.Close()
-	inspectRun(cmd.Context(), d, InspectFlags.URL)
-}
-
-func inspectRun(ctx context.Context, client *sqlclient.Client, url string) {
+	defer client.Close()
 	schemas := InspectFlags.Schema
-	n, err := SchemaNameFromURL(ctx, url)
-	cobra.CheckErr(err)
-	if n != "" {
-		schemas = append(schemas, n)
+	if client.URL.Schema != "" {
+		schemas = append(schemas, client.URL.Schema)
 	}
-	s, err := client.InspectRealm(ctx, &schema.InspectRealmOption{
+	s, err := client.InspectRealm(cmd.Context(), &schema.InspectRealmOption{
 		Schemas: schemas,
 	})
 	cobra.CheckErr(err)
diff --git a/cmd/action/migrate.go b/cmd/action/migrate.go
index 038046c5..a52f40a5 100644
--- a/cmd/action/migrate.go
+++ b/cmd/action/migrate.go
@@ -204,12 +204,6 @@ func to(ctx context.Context, client *sqlclient.Client) (migrate.StateReader, err
 		return nil, fmt.Errorf("invalid driver url %q", MigrateFlags.ToURL)
 	}
 	schemas := MigrateFlags.Schemas
-	if n, err := SchemaNameFromURL(ctx, parts[0]); n != "" {
-		if err != nil {
-			return nil, err
-		}
-		schemas = append(schemas, n)
-	}
 	switch parts[0] {
 	case "file": // hcl file
 		f, err := ioutil.ReadFile(parts[1])
@@ -244,6 +238,9 @@ func to(ctx context.Context, client *sqlclient.Client) (migrate.StateReader, err
 		if err != nil {
 			return nil, err
 		}
+		if client.URL.Schema != "" {
+			schemas = append(schemas, client.URL.Schema)
+		}
 		return struct {
 			migrate.StateReader
 			io.Closer
diff --git a/cmd/action/mux.go b/cmd/action/mux.go
deleted file mode 100644
index 5ce6d668..00000000
--- a/cmd/action/mux.go
+++ /dev/null
@@ -1,144 +0,0 @@
-// Copyright 2021-present The Atlas Authors. All rights reserved.
-// This source code is licensed under the Apache 2.0 license found
-// in the LICENSE file in the root directory of this source tree.
-
-package action
-
-import (
-	"context"
-	"database/sql"
-	"fmt"
-	"net/url"
-	"os"
-	"path/filepath"
-	"regexp"
-	"strings"
-
-	"ariga.io/atlas/sql/sqlite"
-
-	"github.com/go-sql-driver/mysql"
-)
-
-var (
-	reMemMode = regexp.MustCompile(":memory:|^file:.*mode=memory.*")
-)
-
-func urlParts(url string) (string, string, error) {
-	a := strings.SplitN(url, "://", 2)
-	if len(a) != 2 {
-		return "", "", fmt.Errorf(`failed to parse url: "%s"`, url)
-	}
-	return a[0], a[1], nil
-}
-
-// SchemaNameFromURL parses the url the returns schema name
-func SchemaNameFromURL(ctx context.Context, url string) (string, error) {
-	key, dsn, err := urlParts(url)
-	if err != nil {
-		return "", err
-	}
-	switch key {
-	case "mysql", "maria", "mariadb":
-		cfg, err := mysqlConfig(dsn)
-		if err != nil {
-			return "", err
-		}
-		return cfg.DBName, err
-	case "postgres":
-		return postgresSchema(url)
-	case "sqlite":
-		return schemaName(ctx, dsn)
-	default:
-		return "", fmt.Errorf("unknown database type: %q", key)
-	}
-}
-
-func postgresSchema(dsn string) (string, error) {
-	u, err := url.Parse(dsn)
-	if err != nil {
-		return "", err
-	}
-	// lib/pq supports setting default schemas via the `search_path` parameter
-	// in a url.
-	//
-	// See: https://github.com/lib/pq/blob/8446d16b8935fdf2b5c0fe333538ac395e3e1e4b/conn.go#L1155-L1165
-	if schema := u.Query().Get("search_path"); schema != "" {
-		return schema, nil
-	}
-	return "", nil
-}
-
-func schemaName(ctx context.Context, dsn string) (string, error) {
-	err := sqliteFileExists(dsn)
-	if err != nil {
-		return "", err
-	}
-	db, err := sql.Open("sqlite3", dsn)
-	if err != nil {
-		return "", err
-	}
-	drv, err := sqlite.Open(db)
-	if err != nil {
-		return "", err
-	}
-	r, err := drv.InspectRealm(ctx, nil)
-	if err != nil {
-		return "", err
-	}
-	if len(r.Schemas) != 1 {
-		return "", fmt.Errorf("must have exactly 1 schema, got: %d", len(r.Schemas))
-	}
-	return r.Schemas[0].Name, nil
-}
-
-func sqliteFileExists(dsn string) error {
-	if !reMemMode.MatchString(dsn) {
-		return fileExists(dsn)
-	}
-	return nil
-}
-
-func fileExists(dsn string) error {
-	s := strings.Split(dsn, "?")
-	f := dsn
-	if len(s) == 2 {
-		f = s[0]
-	}
-	if strings.Contains(f, "file:") {
-		f = strings.SplitAfter(f, "file:")[1]
-	}
-	f = filepath.Clean(f)
-	if _, err := os.Stat(f); err != nil {
-		return fmt.Errorf("failed opening %q: %w", f, err)
-	}
-	return nil
-}
-
-func mysqlConfig(d string) (*mysql.Config, error) {
-	cfg, err := mysql.ParseDSN(d)
-	// A standard MySQL DSN.
-	if err == nil {
-		return cfg, nil
-	}
-	u, err := url.Parse("mysql://" + d)
-	if err != nil {
-		return nil, err
-	}
-	schema := strings.TrimPrefix(u.Path, "/")
-	// In case of a URL (non-standard DSN),
-	// parse the options from query string.
-	if u.RawQuery != "" {
-		cfg, err = mysql.ParseDSN(fmt.Sprintf("/%s?%s", schema, u.RawQuery))
-		if err != nil {
-			return nil, err
-		}
-	} else {
-		cfg = mysql.NewConfig()
-	}
-	cfg.Net = "tcp"
-	cfg.Addr = u.Host
-	cfg.User = u.User.Username()
-	cfg.Passwd, _ = u.User.Password()
-	cfg.DBName = schema
-	return cfg, nil
-}
diff --git a/cmd/action/version.go b/cmd/action/version.go
index 27223fdd..ab44cbed 100644
--- a/cmd/action/version.go
+++ b/cmd/action/version.go
@@ -21,7 +21,7 @@ var versionCmd = &cobra.Command{
 	Short: "Prints this Atlas CLI version information.",
 	Run: func(cmd *cobra.Command, args []string) {
 		v, u := parse(version)
-		RootCmd.Println(fmt.Sprintf("atlas CLI version %s\n%s", v, u))
+		RootCmd.Printf("atlas CLI version %s\n%s\n", v, u)
 	},
 }
 
diff --git a/sql/mysql/driver.go b/sql/mysql/driver.go
index dba54036..ce879be3 100644
--- a/sql/mysql/driver.go
+++ b/sql/mysql/driver.go
@@ -42,9 +42,12 @@ type (
 func init() {
 	sqlclient.Register(
 		"mysql",
-		sqlclient.DriverOpener(Open, dsn),
+		sqlclient.DriverOpener(Open),
 		sqlclient.RegisterCodec(MarshalHCL, EvalHCL),
 		sqlclient.RegisterFlavours("maria", "mariadb"),
+		sqlclient.RegisterURLParser(func(u *url.URL) *sqlclient.URL {
+			return &sqlclient.URL{URL: u, DSN: dsn(u), Schema: strings.TrimPrefix(u.Path, "/")}
+		}),
 	)
 }
 
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index 23646eb3..f4d50410 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -43,10 +43,11 @@ type (
 func init() {
 	sqlclient.Register(
 		"postgres",
-		sqlclient.DriverOpener(Open, func(u *url.URL) string {
-			return u.String()
-		}),
+		sqlclient.DriverOpener(Open),
 		sqlclient.RegisterCodec(MarshalHCL, EvalHCL),
+		sqlclient.RegisterURLParser(func(u *url.URL) *sqlclient.URL {
+			return &sqlclient.URL{URL: u, DSN: u.String(), Schema: u.Query().Get("search_path")}
+		}),
 	)
 }
 
diff --git a/sql/sqlclient/client.go b/sql/sqlclient/client.go
index c2c92dab..7e01ea17 100644
--- a/sql/sqlclient/client.go
+++ b/sql/sqlclient/client.go
@@ -19,21 +19,37 @@ import (
 	"ariga.io/atlas/sql/migrate"
 )
 
-// Client provides the common functionalities for working with Atlas from different
-// applications (e.g. CLI and TF). Note, the Client is dialect specific and should
-// be instantiated using a call to Open.
-type Client struct {
-	// DB used for creating the client.
-	DB *sql.DB
-
-	// A migration driver for the attached dialect.
-	migrate.Driver
-
-	// Marshal and Unmarshal functions for decoding
-	// and encoding the schema documents.
-	schemaspec.Marshaler
-	schemahcl.Evaluator
-}
+type (
+	// Client provides the common functionalities for working with Atlas from different
+	// applications (e.g. CLI and TF). Note, the Client is dialect specific and should
+	// be instantiated using a call to Open.
+	Client struct {
+		// DB used for creating the client.
+		DB *sql.DB
+		// URL holds an enriched url.URL.
+		URL *URL
+
+		// A migration driver for the attached dialect.
+		migrate.Driver
+
+		// Marshal and Evaluator functions for decoding
+		// and encoding the schema documents.
+		schemaspec.Marshaler
+		schemahcl.Evaluator
+	}
+
+	// URL extends the standard url.URL with additional
+	// connection information attached by the Opener (if any).
+	URL struct {
+		*url.URL
+
+		// The DSN used for opening the connection.
+		DSN string
+
+		// The Schema this client is connected to.
+		Schema string
+	}
+)
 
 // Close closes the underlying database connection and the migration
 // driver in case it implements the io.Closer interface.
@@ -59,9 +75,10 @@ type (
 	// OpenerFunc allows using a function as an Opener.
 	OpenerFunc func(context.Context, *url.URL) (*Client, error)
 
-	namedOpener struct {
+	driver struct {
 		Opener
-		name string
+		name  string
+		parse func(*url.URL) *URL
 	}
 )
 
@@ -82,11 +99,19 @@ func Open(ctx context.Context, s string) (*Client, error) {
 	if !ok {
 		return nil, fmt.Errorf("sql/sqlclient: no opener was register with name %q", u.Scheme)
 	}
-	return v.(namedOpener).Open(ctx, u)
+	client, err := v.(driver).Open(ctx, u)
+	if err != nil {
+		return nil, err
+	}
+	if client.URL == nil {
+		client.URL = v.(driver).parse(u)
+	}
+	return client, nil
 }
 
 type (
 	registerOptions struct {
+		parse    func(*url.URL) *URL
 		flavours []string
 		codec    interface {
 			schemaspec.Marshaler
@@ -106,6 +131,14 @@ func RegisterFlavours(flavours ...string) RegisterOption {
 	}
 }
 
+// RegisterURLParser allows registering a function for parsing
+// the url.URL and attach additional info to the extended URL.
+func RegisterURLParser(p func(*url.URL) *URL) RegisterOption {
+	return func(opts *registerOptions) {
+		opts.parse = p
+	}
+}
+
 // RegisterCodec registers static codec for attaching into
 // the client after it is opened.
 func RegisterCodec(m schemaspec.Marshaler, e schemahcl.Evaluator) RegisterOption {
@@ -121,13 +154,14 @@ func RegisterCodec(m schemaspec.Marshaler, e schemahcl.Evaluator) RegisterOption
 }
 
 // DriverOpener is a helper Opener creator for sharing between all drivers.
-func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error), dsn func(*url.URL) string) Opener {
+func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error)) Opener {
 	return OpenerFunc(func(ctx context.Context, u *url.URL) (*Client, error) {
 		v, ok := drivers.Load(u.Scheme)
 		if !ok {
 			return nil, fmt.Errorf("sql/sqlclient: unexpected missing opener %q", u.Scheme)
 		}
-		db, err := sql.Open(v.(namedOpener).name, dsn(u))
+		ur := v.(driver).parse(u)
+		db, err := sql.Open(v.(driver).name, ur.DSN)
 		if err != nil {
 			return nil, err
 		}
@@ -140,6 +174,7 @@ func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error), dsn fun
 		}
 		return &Client{
 			DB:     db,
+			URL:    ur,
 			Driver: drv,
 		}, nil
 	})
@@ -150,7 +185,10 @@ func Register(name string, opener Opener, opts ...RegisterOption) {
 	if opener == nil {
 		panic("sql/sqlclient: Register opener is nil")
 	}
-	opt := &registerOptions{}
+	opt := &registerOptions{
+		// Default URL parser uses the URL as the DSN.
+		parse: func(u *url.URL) *URL { return &URL{URL: u, DSN: u.String()} },
+	}
 	for i := range opts {
 		opts[i](opt)
 	}
@@ -169,8 +207,9 @@ func Register(name string, opener Opener, opts ...RegisterOption) {
 		if _, ok := drivers.Load(f); ok {
 			panic("sql/sqlclient: Register called twice for " + f)
 		}
-		drivers.Store(f, namedOpener{
+		drivers.Store(f, driver{
 			name:   name,
+			parse:  opt.parse,
 			Opener: opener,
 		})
 	}
diff --git a/sql/sqlite/driver.go b/sql/sqlite/driver.go
index 76274d58..0a862111 100644
--- a/sql/sqlite/driver.go
+++ b/sql/sqlite/driver.go
@@ -39,11 +39,12 @@ type (
 func init() {
 	sqlclient.Register(
 		"sqlite3",
-		sqlclient.DriverOpener(Open, func(u *url.URL) string {
-			return strings.TrimPrefix(u.String(), u.Scheme+"://")
-		}),
+		sqlclient.DriverOpener(Open),
 		sqlclient.RegisterCodec(MarshalHCL, EvalHCL),
 		sqlclient.RegisterFlavours("sqlite"),
+		sqlclient.RegisterURLParser(func(u *url.URL) *sqlclient.URL {
+			return &sqlclient.URL{URL: u, DSN: strings.TrimPrefix(u.String(), u.Scheme+"://"), Schema: mainFile}
+		}),
 	)
 }
 
@@ -80,7 +81,7 @@ func (d *Driver) IsClean(ctx context.Context) (bool, error) {
 	if err != nil {
 		return false, err
 	}
-	return r == nil || len(r.Schemas) == 1 && r.Schemas[0].Name == "main" && len(r.Schemas[0].Tables) == 0, nil
+	return r == nil || len(r.Schemas) == 1 && r.Schemas[0].Name == mainFile && len(r.Schemas[0].Tables) == 0, nil
 }
 
 // Clean implements the inlined migrate.Clean interface to override the "emptying" behavior.


Test patch:
diff --git a/cmd/action/mux_test.go b/cmd/action/mux_test.go
deleted file mode 100644
index 02932d8e..00000000
--- a/cmd/action/mux_test.go
+++ /dev/null
@@ -1,127 +0,0 @@
-// Copyright 2021-present The Atlas Authors. All rights reserved.
-// This source code is licensed under the Apache 2.0 license found
-// in the LICENSE file in the root directory of this source tree.
-
-package action_test
-
-import (
-	"context"
-	"io/ioutil"
-	"os"
-	"testing"
-
-	"ariga.io/atlas/cmd/action"
-
-	_ "github.com/mattn/go-sqlite3"
-	"github.com/stretchr/testify/require"
-)
-
-func Test_SQLiteFileDoestNotExist(t *testing.T) {
-	var tests = []struct {
-		url      string
-		expected string
-	}{
-		{
-			url:      "sqlite://test.db",
-			expected: `failed opening "test.db": stat test.db: no such file or directory`,
-		},
-		{
-			url:      "sqlite://some_random_string_like_this",
-			expected: `failed opening "some_random_string_like_this": stat some_random_string_like_this: no such file or directory`,
-		},
-		{
-			url:      "sqlite://file:/home/fred/data.db",
-			expected: `failed opening "/home/fred/data.db": stat /home/fred/data.db: no such file or directory`,
-		},
-		{
-			url:      "sqlite://file:///home/fred/data.db",
-			expected: `failed opening "/home/fred/data.db": stat /home/fred/data.db: no such file or directory`,
-		},
-		{
-			url:      "sqlite://file://localhost/home/fred/data.db",
-			expected: `failed opening "/localhost/home/fred/data.db": stat /localhost/home/fred/data.db: no such file or directory`,
-		},
-		{
-			url:      "sqlite://file://darkstar/home/fred/data.db",
-			expected: `failed opening "/darkstar/home/fred/data.db": stat /darkstar/home/fred/data.db: no such file or directory`,
-		},
-		{
-			url:      "sqlite://file:data.db?mode=ro&cache=private",
-			expected: `failed opening "data.db": stat data.db: no such file or directory`,
-		},
-		{
-			url:      "sqlite://file:/home/fred/data.db?vfs=unix-dotfile",
-			expected: `failed opening "/home/fred/data.db": stat /home/fred/data.db: no such file or directory`,
-		},
-		{
-			url:      "sqlite://file:data.db?mode=readonly",
-			expected: `failed opening "data.db": stat data.db: no such file or directory`,
-		},
-		{
-			url:      "sqlite://asdad?cache=shared&mode=memory",
-			expected: `failed opening "asdad": stat asdad: no such file or directory`,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.url, func(t *testing.T) {
-			_, err := action.SchemaNameFromURL(context.Background(), tt.url)
-			require.EqualError(t, err, tt.expected)
-		})
-	}
-}
-
-func Test_SQLiteFileExist(t *testing.T) {
-	r := require.New(t)
-	file, err := ioutil.TempFile("", "tmp")
-	r.NoError(err)
-	t.Cleanup(func() {
-		err := os.Remove(file.Name())
-		r.NoError(err)
-	})
-	dsn := "sqlite://file://" + file.Name()
-	_, err = action.SchemaNameFromURL(context.Background(), dsn)
-	r.NoError(err)
-}
-
-func Test_SQLiteInMemory(t *testing.T) {
-	r := require.New(t)
-	_, err := action.SchemaNameFromURL(context.Background(), "sqlite://file:test.db?cache=shared&mode=memory")
-	r.NoError(err)
-}
-
-func Test_PostgresSchemaDSN(t *testing.T) {
-	var tests = []struct {
-		url      string
-		expected string
-		wantErr  bool
-	}{
-		{
-			url:      "postgres://localhost:5432/dbname?search_path=foo",
-			expected: "foo",
-		},
-		{
-			url:      "postgres://localhost:5432/dbname",
-			expected: "",
-		},
-		{
-			url:      "postgres://(bad:host)?search_path=foo",
-			expected: "",
-			wantErr:  true,
-		},
-		{
-			url:      "postgres://localhost:5432/dbname?search_path=",
-			expected: "",
-		},
-		{
-			url:      "postgres://user_name:password@localhost:5432/dbname?sslmode=disable",
-			expected: "",
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.url, func(t *testing.T) {
-			schema, err := action.SchemaNameFromURL(context.Background(), tt.url)
-			require.Equal(t, tt.wantErr, err != nil)
-			require.Equal(t, tt.expected, schema)
-		})
-	}
-}
diff --git a/sql/sqlclient/client_test.go b/sql/sqlclient/client_test.go
index a352c74e..f38cda13 100644
--- a/sql/sqlclient/client_test.go
+++ b/sql/sqlclient/client_test.go
@@ -22,6 +22,9 @@ func TestRegisterOpen(t *testing.T) {
 			return c, nil
 		}),
 		sqlclient.RegisterFlavours("maria"),
+		sqlclient.RegisterURLParser(func(u *url.URL) *sqlclient.URL {
+			return &sqlclient.URL{URL: u, DSN: "dsn", Schema: "schema"}
+		}),
 	)
 	require.PanicsWithValue(
 		t,
@@ -40,9 +43,15 @@ func TestRegisterOpen(t *testing.T) {
 	c1, err := sqlclient.Open(context.Background(), "mysql://:3306")
 	require.NoError(t, err)
 	require.True(t, c == c1)
+	require.Equal(t, "dsn", c.URL.DSN)
+	require.Equal(t, "schema", c.URL.Schema)
+
 	c1, err = sqlclient.Open(context.Background(), "maria://:3306")
 	require.NoError(t, err)
 	require.True(t, c == c1)
+	require.Equal(t, "dsn", c.URL.DSN)
+	require.Equal(t, "schema", c.URL.Schema)
+
 	c1, err = sqlclient.Open(context.Background(), "postgres://:3306")
 	require.EqualError(t, err, `sql/sqlclient: no opener was register with name "postgres"`)
 }