Repository: ariga/atlas
Pull request #1162 (2022-09-25 10:10:16)

sql/sqlite: toggle foreign-keys pragma before opening transactions


Patch:
diff --git a/sql/sqlclient/client.go b/sql/sqlclient/client.go
index 9d587321..29105092 100644
--- a/sql/sqlclient/client.go
+++ b/sql/sqlclient/client.go
@@ -42,8 +42,9 @@ type (
 		schemahcl.Marshaler
 		schemahcl.Evaluator
 
-		// A func to open a migrate.Driver with a given schema.ExecQuerier. Used when creating a TxClient.
+		// Functions registered by the drivers and used for opening transactions and their clients.
 		openDriver func(schema.ExecQuerier) (migrate.Driver, error)
+		openTx     TxOpener
 	}
 
 	// TxClient is returned by calling Client.Tx. It behaves the same as Client,
@@ -52,7 +53,7 @@ type (
 		*Client
 
 		// The transaction this Client wraps.
-		Tx *sql.Tx
+		Tx *Tx
 	}
 
 	// URL extends the standard url.URL with additional
@@ -73,9 +74,20 @@ func (c *Client) Tx(ctx context.Context, opts *sql.TxOptions) (*TxClient, error)
 	if c.openDriver == nil {
 		return nil, errors.New("sql/sqlclient: unexpected driver opener: <nil>")
 	}
-	tx, err := c.DB.BeginTx(ctx, opts)
-	if err != nil {
-		return nil, fmt.Errorf("sql/sqlclient: starting transaction: %w", err)
+	var tx *Tx
+	switch {
+	case c.openTx != nil:
+		ttx, err := c.openTx(ctx, c.DB, opts)
+		if err != nil {
+			return nil, err
+		}
+		tx = ttx
+	default:
+		ttx, err := c.DB.BeginTx(ctx, opts)
+		if err != nil {
+			return nil, fmt.Errorf("sql/sqlclient: starting transaction: %w", err)
+		}
+		tx = &Tx{Tx: ttx}
 	}
 	drv, err := c.openDriver(tx)
 	if err != nil {
@@ -88,12 +100,27 @@ func (c *Client) Tx(ctx context.Context, opts *sql.TxOptions) (*TxClient, error)
 
 // Commit the transaction.
 func (c *TxClient) Commit() error {
-	return c.Tx.Commit()
+	if err := c.Tx.Commit(); err != nil {
+		return err
+	}
+	return c.close()
 }
 
 // Rollback the transaction.
 func (c *TxClient) Rollback() error {
-	return c.Tx.Rollback()
+	if err := c.Tx.Rollback(); err != nil {
+		return err
+	}
+	return c.close()
+}
+
+func (c *TxClient) close() error {
+	if c.Tx.Close != nil {
+		if err := c.Tx.Close(); err != nil {
+			return fmt.Errorf("sql/sqlclient: closing transaction: %w", err)
+		}
+	}
+	return nil
 }
 
 // AddClosers adds list of closers to close at the end of the client lifetime.
@@ -139,8 +166,9 @@ type (
 
 	driver struct {
 		Opener
-		name   string
-		parser URLParser
+		name     string
+		parser   URLParser
+		txOpener TxOpener
 	}
 )
 
@@ -220,6 +248,7 @@ func OpenSchema(s string) OpenOption {
 type (
 	registerOptions struct {
 		openDriver func(schema.ExecQuerier) (migrate.Driver, error)
+		txOpener   TxOpener
 		parser     URLParser
 		flavours   []string
 		codec      interface {
@@ -278,12 +307,13 @@ func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error)) Opener
 		if !ok {
 			return nil, fmt.Errorf("sql/sqlclient: unexpected missing opener %q", u.Scheme)
 		}
-		ur := v.(*driver).parser.ParseURL(u)
-		db, err := sql.Open(v.(*driver).name, ur.DSN)
+		drv := v.(*driver)
+		ur := drv.parser.ParseURL(u)
+		db, err := sql.Open(drv.name, ur.DSN)
 		if err != nil {
 			return nil, err
 		}
-		drv, err := open(db)
+		mdr, err := open(db)
 		if err != nil {
 			if cerr := db.Close(); cerr != nil {
 				err = fmt.Errorf("%w: %v", err, cerr)
@@ -291,15 +321,33 @@ func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error)) Opener
 			return nil, err
 		}
 		return &Client{
-			Name:       v.(*driver).name,
+			Name:       drv.name,
 			DB:         db,
 			URL:        ur,
-			Driver:     drv,
+			Driver:     mdr,
 			openDriver: open,
+			openTx:     drv.txOpener,
 		}, nil
 	})
 }
 
+type (
+	// Tx wraps sql.Tx with an optional closer.
+	Tx struct {
+		*sql.Tx
+		Close func() error
+	}
+	// TxOpener opens a transaction with optional closer.
+	TxOpener func(context.Context, *sql.DB, *sql.TxOptions) (*Tx, error)
+)
+
+// RegisterTxOpener allows registering a custom transaction opener with an optional close function.
+func RegisterTxOpener(open TxOpener) RegisterOption {
+	return func(opts *registerOptions) {
+		opts.txOpener = open
+	}
+}
+
 // Register registers a client Opener (i.e. creator) with the given name.
 func Register(name string, opener Opener, opts ...RegisterOption) {
 	if opener == nil {
@@ -335,7 +383,7 @@ func Register(name string, opener Opener, opts ...RegisterOption) {
 			return c, err
 		})
 	}
-	drv := &driver{opener, name, opt.parser}
+	drv := &driver{opener, name, opt.parser, opt.txOpener}
 	for _, f := range append(opt.flavours, name) {
 		if _, ok := drivers.Load(f); ok {
 			panic("sql/sqlclient: Register called twice for " + f)
diff --git a/sql/sqlite/driver.go b/sql/sqlite/driver.go
index 06844b1f..93edd8ea 100644
--- a/sql/sqlite/driver.go
+++ b/sql/sqlite/driver.go
@@ -6,6 +6,7 @@ package sqlite
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 	"fmt"
 	"net/url"
@@ -35,7 +36,6 @@ type (
 	conn struct {
 		schema.ExecQuerier
 		// System variables that are set on `Open`.
-		fkEnabled  bool
 		version    string
 		collations []string
 	}
@@ -48,6 +48,7 @@ func init() {
 	sqlclient.Register(
 		DriverName,
 		sqlclient.DriverOpener(Open),
+		sqlclient.RegisterTxOpener(openTx),
 		sqlclient.RegisterCodec(MarshalHCL, EvalHCL),
 		sqlclient.RegisterFlavours("sqlite"),
 		sqlclient.RegisterURLParser(sqlclient.URLParserFunc(func(u *url.URL) *sqlclient.URL {
@@ -67,11 +68,11 @@ func Open(db schema.ExecQuerier) (migrate.Driver, error) {
 		c   = conn{ExecQuerier: db}
 		ctx = context.Background()
 	)
-	rows, err := db.QueryContext(ctx, "SELECT sqlite_version(), foreign_keys from pragma_foreign_keys")
+	rows, err := db.QueryContext(ctx, "SELECT sqlite_version()")
 	if err != nil {
-		return nil, fmt.Errorf("sqlite: query version and foreign_keys pragma: %w", err)
+		return nil, fmt.Errorf("sqlite: query version pragma: %w", err)
 	}
-	if err := sqlx.ScanOne(rows, &c.version, &c.fkEnabled); err != nil {
+	if err := sqlx.ScanOne(rows, &c.version); err != nil {
 		return nil, fmt.Errorf("sqlite: scan version and foreign_keys pragma: %w", err)
 	}
 	if rows, err = db.QueryContext(ctx, "SELECT name FROM pragma_collation_list()"); err != nil {
@@ -164,6 +165,35 @@ func acquireLock(path string, timeout time.Duration) (schema.UnlockFunc, error)
 	return func() error { return os.Remove(path) }, nil
 }
 
+func openTx(ctx context.Context, db *sql.DB, opts *sql.TxOptions) (*sqlclient.Tx, error) {
+	var on sql.NullBool
+	if err := db.QueryRowContext(ctx, "PRAGMA foreign_keys").Scan(&on); err != nil {
+		return nil, fmt.Errorf("sql/sqlite: querying 'foreign_keys' pragma: %w", err)
+	}
+	// Disable the foreign_keys pragma in case it is enabled, and
+	// toggle it back after transaction is committed or rolled back.
+	if on.Bool {
+		if _, err := db.ExecContext(ctx, "PRAGMA foreign_keys = off"); err != nil {
+			return nil, fmt.Errorf("sql/sqlite: set 'foreign_keys = off': %w", err)
+		}
+	}
+	tx, err := db.BeginTx(ctx, opts)
+	if err != nil {
+		return nil, err
+	}
+	return &sqlclient.Tx{
+		Tx: tx,
+		Close: func() error {
+			if on.Bool {
+				if _, err := db.ExecContext(ctx, "PRAGMA foreign_keys = on"); err != nil {
+					return fmt.Errorf("sql/sqlite: set 'foreign_keys = on': %w", err)
+				}
+			}
+			return nil
+		},
+	}, nil
+}
+
 // SQLite standard data types as defined in its codebase and documentation.
 // https://www.sqlite.org/datatype3.html
 // https://github.com/sqlite/sqlite/blob/master/src/global.c
diff --git a/sql/sqlite/migrate.go b/sql/sqlite/migrate.go
index e6fbea93..0b2dca48 100644
--- a/sql/sqlite/migrate.go
+++ b/sql/sqlite/migrate.go
@@ -85,9 +85,9 @@ func (s *state) plan(ctx context.Context, changes []schema.Change) (err error) {
 	}
 	// Disable foreign-keys enforcement if it is required
 	// by one of the changes in the plan.
-	if s.skipFKs && s.conn.fkEnabled {
+	if s.skipFKs {
 		// Callers should note that these 2 pragmas are no-op in transactions,
-		// and therefore, should not call BEGIN manually. https://sqlite.org/pragma.html#pragma_foreign_keys
+		// See: https://sqlite.org/pragma.html#pragma_foreign_keys.
 		s.Changes = append([]*migrate.Change{{Cmd: "PRAGMA foreign_keys = off", Comment: "disable the enforcement of foreign-keys constraints"}}, s.Changes...)
 		s.append(&migrate.Change{Cmd: "PRAGMA foreign_keys = on", Comment: "enable back the enforcement of foreign-keys constraints"})
 	}


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate_test.go b/cmd/atlas/internal/cmdapi/migrate_test.go
index 94f58ae9..10f5ad46 100644
--- a/cmd/atlas/internal/cmdapi/migrate_test.go
+++ b/cmd/atlas/internal/cmdapi/migrate_test.go
@@ -6,6 +6,7 @@ package cmdapi
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 	"fmt"
 	"io"
@@ -367,6 +368,39 @@ func TestMigrate_Apply(t *testing.T) {
 	require.EqualError(t, err, migrate.MissingMigrationError{Version: "z", Description: "z"}.Error())
 }
 
+func TestMigrate_ApplyTxMode(t *testing.T) {
+	for _, mode := range []string{"none", "file", "all"} {
+		t.Run(mode, func(t *testing.T) {
+			p := t.TempDir()
+			// Apply the first 2 migrations.
+			s, err := runCmd(
+				Root, "migrate", "apply",
+				"--dir", "file://testdata/sqlitetx",
+				"--url", fmt.Sprintf("sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test.db")),
+				"--tx-mode", mode,
+				"2",
+			)
+			require.NoError(t, err)
+			require.NotEmpty(t, s)
+			db, err := sql.Open("sqlite3", fmt.Sprintf("file:%s?cache=shared&_fk=1", filepath.Join(p, "test.db")))
+			require.NoError(t, err)
+			var n int
+			require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM `friendships`").Scan(&n))
+			require.Equal(t, 2, n)
+
+			// Apply the rest.
+			s, err = runCmd(
+				Root, "migrate", "apply",
+				"--dir", "file://testdata/sqlitetx",
+				"--url", fmt.Sprintf("sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test.db")),
+				"--tx-mode", mode,
+			)
+			require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM `friendships`").Scan(&n))
+			require.Equal(t, 2, n)
+		})
+	}
+}
+
 func TestMigrate_ApplyBaseline(t *testing.T) {
 	p := t.TempDir()
 	// Run migration with baseline should store this revision in the database.
diff --git a/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925092817_initial.sql b/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925092817_initial.sql
new file mode 100644
index 00000000..ffd24610
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925092817_initial.sql
@@ -0,0 +1,2 @@
+-- create "users" table
+CREATE TABLE `users` (`id` integer NOT NULL, `name` text NULL, PRIMARY KEY (`id`));
diff --git a/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925094021_second.sql b/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925094021_second.sql
new file mode 100644
index 00000000..b2c69b70
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925094021_second.sql
@@ -0,0 +1,5 @@
+-- create "friendships" table
+CREATE TABLE `friendships` (`user_id` integer NOT NULL, `friend_id` integer NOT NULL, PRIMARY KEY (`user_id`, `friend_id`), CONSTRAINT `user_id_fk` FOREIGN KEY (`user_id`) REFERENCES `users` (`id`) ON DELETE CASCADE, CONSTRAINT `friend_id_fk` FOREIGN KEY (`friend_id`) REFERENCES `users` (`id`) ON DELETE CASCADE);
+
+INSERT INTO `users` (`id`) VALUES (1), (2);
+INSERT INTO `friendships` (`user_id`, `friend_id`) VALUES (1,2), (2,1);
diff --git a/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925094437_third.sql b/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925094437_third.sql
new file mode 100644
index 00000000..6544f7fa
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/sqlitetx/20220925094437_third.sql
@@ -0,0 +1,12 @@
+-- disable the enforcement of foreign-keys constraints
+PRAGMA foreign_keys = off;
+-- create "new_users" table
+CREATE TABLE `new_users` (`id` integer NOT NULL, PRIMARY KEY (`id`));
+-- copy rows from old table "users" to new temporary table "new_users"
+INSERT INTO `new_users` (`id`) SELECT `id` FROM `users`;
+-- drop "users" table after copying rows
+DROP TABLE `users`;
+-- rename temporary table "new_users" to "users"
+ALTER TABLE `new_users` RENAME TO `users`;
+-- enable back the enforcement of foreign-keys constraints
+PRAGMA foreign_keys = on;
diff --git a/cmd/atlas/internal/cmdapi/testdata/sqlitetx/atlas.sum b/cmd/atlas/internal/cmdapi/testdata/sqlitetx/atlas.sum
new file mode 100644
index 00000000..06f69682
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/sqlitetx/atlas.sum
@@ -0,0 +1,4 @@
+h1:09lkmQGTxdoqPwK0ZXtU4+KHipufkVp1Jlje3t6Opy4=
+20220925092817_initial.sql h1:ZGeLdeqNUMXqJm+hPkhBrhzbtUzSBH8yVTsnSnJo/qU=
+20220925094021_second.sql h1:vcoquz3yk+TlTPiQgW5hHpS/abIvySCM/bzgwYTDoqY=
+20220925094437_third.sql h1:2pbBiUBKsEC5+ppfPPTDr+iwJSgZ2rM4qmHI44/vmnc=
diff --git a/sql/sqlclient/client_test.go b/sql/sqlclient/client_test.go
index 3ea67565..37a4b293 100644
--- a/sql/sqlclient/client_test.go
+++ b/sql/sqlclient/client_test.go
@@ -96,6 +96,11 @@ func TestClient_Tx(t *testing.T) {
 		sqlclient.RegisterDriverOpener(func(db schema.ExecQuerier) (migrate.Driver, error) {
 			return &mockDriver{db: db}, nil
 		}),
+		sqlclient.RegisterTxOpener(func(ctx context.Context, db *sql.DB, opts *sql.TxOptions) (*sqlclient.Tx, error) {
+			tx, err := db.BeginTx(ctx, opts)
+			require.NoError(t, err)
+			return &sqlclient.Tx{Tx: tx, Close: func() error { return nil }}, nil
+		}),
 	)
 
 	c, err := sqlclient.Open(context.Background(), "tx://")
@@ -107,6 +112,7 @@ func TestClient_Tx(t *testing.T) {
 	_, err = tx.ExecContext(context.Background(), stmt)
 	require.NoError(t, err)
 	require.NoError(t, tx.Commit())
+	require.NotNil(t, tx.Close)
 
 	// Rollback works as well.
 	tx, err = c.Tx(context.Background(), nil)
diff --git a/sql/sqlite/inspect_test.go b/sql/sqlite/inspect_test.go
index 09f5f7e6..81477a21 100644
--- a/sql/sqlite/inspect_test.go
+++ b/sql/sqlite/inspect_test.go
@@ -501,11 +501,11 @@ type mock struct {
 }
 
 func (m mock) systemVars(version string) {
-	m.ExpectQuery(sqltest.Escape("SELECT sqlite_version(), foreign_keys from pragma_foreign_keys")).
+	m.ExpectQuery(sqltest.Escape("SELECT sqlite_version()")).
 		WillReturnRows(sqltest.Rows(`
-     version    |   foreign_keys    
-----------------+-----------------
- ` + version + `|       1
+     version    
+----------------
+ ` + version + `
 `))
 	m.ExpectQuery(sqltest.Escape("SELECT name FROM pragma_collation_list()")).
 		WillReturnRows(sqltest.Rows(`