Repository: ariga/atlas
Pull request #3287 (2024-12-28 06:51:32)

sql/sqlite: lock by database url
Approved internally. 


Patch:
diff --git a/sql/sqlite/driver.go b/sql/sqlite/driver.go
index 0a9d6858..f98cd09f 100644
--- a/sql/sqlite/driver.go
+++ b/sql/sqlite/driver.go
@@ -9,6 +9,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"hash/adler32"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -35,9 +36,9 @@ type (
 	// database connection and its information.
 	conn struct {
 		schema.ExecQuerier
+		url *sqlclient.URL
 		// System variables that are set on `Open`.
-		version    string
-		collations []string
+		version string
 	}
 )
 
@@ -54,18 +55,12 @@ const DriverName = "sqlite3"
 func init() {
 	sqlclient.Register(
 		DriverName,
-		sqlclient.DriverOpener(Open),
+		sqlclient.OpenerFunc(opener),
+		sqlclient.RegisterDriverOpener(Open),
 		sqlclient.RegisterTxOpener(OpenTx),
 		sqlclient.RegisterCodec(MarshalHCL, EvalHCL),
 		sqlclient.RegisterFlavours("sqlite"),
-		sqlclient.RegisterURLParser(sqlclient.URLParserFunc(func(u *url.URL) *sqlclient.URL {
-			uc := &sqlclient.URL{URL: u, DSN: strings.TrimPrefix(u.String(), u.Scheme+"://"), Schema: mainFile}
-			if mode := u.Query().Get("mode"); mode == "memory" {
-				// The "file:" prefix is mandatory for memory modes.
-				uc.DSN = "file:" + uc.DSN
-			}
-			return uc
-		})),
+		sqlclient.RegisterURLParser(urlparse{}),
 	)
 	sqlclient.Register(
 		"libsql",
@@ -83,6 +78,42 @@ func init() {
 	)
 }
 
+type urlparse struct{}
+
+// ParseURL implements the sqlclient.URLParser interface.
+func (urlparse) ParseURL(u *url.URL) *sqlclient.URL {
+	uc := &sqlclient.URL{URL: u, DSN: strings.TrimPrefix(u.String(), u.Scheme+"://"), Schema: mainFile}
+	if mode := u.Query().Get("mode"); mode == "memory" {
+		// The "file:" prefix is mandatory for memory modes.
+		uc.DSN = "file:" + uc.DSN
+	}
+	return uc
+}
+
+func opener(_ context.Context, u *url.URL) (*sqlclient.Client, error) {
+	ur := urlparse{}.ParseURL(u)
+	db, err := sql.Open(DriverName, ur.DSN)
+	if err != nil {
+		return nil, err
+	}
+	drv, err := Open(db)
+	if err != nil {
+		if cerr := db.Close(); cerr != nil {
+			err = fmt.Errorf("%w: %v", err, cerr)
+		}
+		return nil, err
+	}
+	if drv, ok := drv.(*Driver); ok {
+		drv.url = ur
+	}
+	return &sqlclient.Client{
+		Name:   DriverName,
+		DB:     db,
+		URL:    ur,
+		Driver: drv,
+	}, nil
+}
+
 // Open opens a new SQLite driver.
 func Open(db schema.ExecQuerier) (migrate.Driver, error) {
 	var (
@@ -96,12 +127,6 @@ func Open(db schema.ExecQuerier) (migrate.Driver, error) {
 	if err := sqlx.ScanOne(rows, &c.version); err != nil {
 		return nil, fmt.Errorf("sqlite: scan version pragma: %w", err)
 	}
-	if rows, err = db.QueryContext(ctx, "SELECT name FROM pragma_collation_list()"); err != nil {
-		return nil, fmt.Errorf("sqlite: query collation_list pragma: %w", err)
-	}
-	if c.collations, err = sqlx.ScanStrings(rows); err != nil {
-		return nil, fmt.Errorf("sqlite: scanning database collations: %w", err)
-	}
 	return &Driver{
 		conn:        c,
 		Differ:      &sqlx.Diff{DiffDriver: &diff{}},
@@ -155,6 +180,11 @@ func (d *Driver) CheckClean(ctx context.Context, revT *migrate.TableIdent) error
 
 // Lock implements the schema.Locker interface.
 func (d *Driver) Lock(_ context.Context, name string, timeout time.Duration) (schema.UnlockFunc, error) {
+	// If the URL was set and the database is a file, use its name in the lock file.
+	if d.url != nil && strings.HasPrefix(d.url.DSN, "file:") {
+		p := filepath.Join(d.url.Host, d.url.Path)
+		name = fmt.Sprintf("%s_%s", name, fmt.Sprintf("%x", adler32.Checksum([]byte(p))))
+	}
 	path := filepath.Join(os.TempDir(), name+".lock")
 	c, err := os.ReadFile(path)
 	if errors.Is(err, os.ErrNotExist) {


Test patch:
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index d2c6e9af..5c734aa4 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -16,6 +16,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"slices"
 	"strings"
 	"sync"
 	"testing"
@@ -451,17 +452,15 @@ func testCLISchemaApply(t T, h string, url string, args ...string) {
 	f := filepath.Join(t.TempDir(), "schema.hcl")
 	err := os.WriteFile(f, []byte(h), 0644)
 	require.NoError(t, err)
-	runArgs := []string{
-		"schema",
-		"apply",
-		"-u",
-		url,
-		"-f",
-		f,
-		"--dev-url",
-		url,
+	runArgs := append([]string{
+		"schema", "apply",
+		"-u", url,
+		"-f", f,
+	}, args...)
+	// Append the --dev-url only if it is not already present.
+	if !slices.Contains(args, "--dev-url") {
+		args = append(args, "--dev-url", url)
 	}
-	runArgs = append(runArgs, args...)
 	cmd := exec.Command(execPath(t), runArgs...)
 	stdout, stderr := bytes.NewBuffer(nil), bytes.NewBuffer(nil)
 	cmd.Stderr = stderr
diff --git a/internal/integration/sqlite_test.go b/internal/integration/sqlite_test.go
index bb746273..c7f69348 100644
--- a/internal/integration/sqlite_test.go
+++ b/internal/integration/sqlite_test.go
@@ -507,7 +507,7 @@ env "hello" {
 	})
 	t.Run("SchemaApply", func(t *testing.T) {
 		liteRun(t, func(t *liteTest) {
-			testCLISchemaApply(t, h, t.url(""))
+			testCLISchemaApply(t, h, t.url(""), "--dev-url", t.dev())
 		})
 	})
 	t.Run("SchemaApplyWithVars", func(t *testing.T) {
@@ -526,7 +526,7 @@ table "users" {
 }
 `
 		liteRun(t, func(t *liteTest) {
-			testCLISchemaApply(t, h, t.url(""), "--var", "tenant=main")
+			testCLISchemaApply(t, h, t.url(""), "--var", "tenant=main", "--dev-url", t.dev())
 		})
 	})
 	t.Run("SchemaApplyDryRun", func(t *testing.T) {
@@ -938,6 +938,10 @@ func (t *liteTest) url(_ string) string {
 	return fmt.Sprintf("sqlite://file:%s?cache=shared&_fk=1", t.file)
 }
 
+func (t *liteTest) dev() string {
+	return fmt.Sprintf("sqlite://%s?mode=memory&_fk=1", t.Name())
+}
+
 func (t *liteTest) applyRealmHcl(spec string) {
 	t.applyHcl(spec)
 }
diff --git a/sql/sqlite/driver_test.go b/sql/sqlite/driver_test.go
index f7d7904b..fbb280f4 100644
--- a/sql/sqlite/driver_test.go
+++ b/sql/sqlite/driver_test.go
@@ -45,7 +45,7 @@ func TestParser_ParseURL(t *testing.T) {
 }
 
 func TestDriver_LockAcquired(t *testing.T) {
-	drv := &Driver{}
+	drv := &Driver{conn: &conn{}}
 
 	// Acquiring a lock does work.
 	unlock, err := drv.Lock(context.Background(), "lock", time.Second)
diff --git a/sql/sqlite/inspect_test.go b/sql/sqlite/inspect_test.go
index 7d1db781..6bdb1606 100644
--- a/sql/sqlite/inspect_test.go
+++ b/sql/sqlite/inspect_test.go
@@ -566,16 +566,6 @@ func (m mock) systemVars(version string) {
      version    
 ----------------
  ` + version + `
-`))
-	m.ExpectQuery(sqltest.Escape("SELECT name FROM pragma_collation_list()")).
-		WillReturnRows(sqltest.Rows(`
-  pragma_collation_list   
-------------------------
-      decimal
-      uint
-      RTRIM
-      NOCASE
-      BINARY
 `))
 }