Repository: ariga/atlas
Pull request #1164 (2022-09-26 13:53:52)

sql/sqlite: check foreign keys once before committing a transaction


Patch:
diff --git a/sql/sqlclient/client.go b/sql/sqlclient/client.go
index 29105092..21a293d3 100644
--- a/sql/sqlclient/client.go
+++ b/sql/sqlclient/client.go
@@ -100,27 +100,12 @@ func (c *Client) Tx(ctx context.Context, opts *sql.TxOptions) (*TxClient, error)
 
 // Commit the transaction.
 func (c *TxClient) Commit() error {
-	if err := c.Tx.Commit(); err != nil {
-		return err
-	}
-	return c.close()
+	return c.Tx.Commit()
 }
 
 // Rollback the transaction.
 func (c *TxClient) Rollback() error {
-	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
+	return c.Tx.Rollback()
 }
 
 // AddClosers adds list of closers to close at the end of the client lifetime.
@@ -218,20 +203,24 @@ func OpenURL(ctx context.Context, u *url.URL, opts ...OpenOption) (*Client, erro
 	if !ok {
 		return nil, fmt.Errorf("sql/sqlclient: no opener was register with name %q", u.Scheme)
 	}
+	drv := v.(*driver)
 	// If there is a schema given and the driver allows to change the schema for the url, do it.
 	if cfg.schema != nil {
-		sc, ok := v.(*driver).parser.(SchemaChanger)
+		sc, ok := drv.parser.(SchemaChanger)
 		if !ok {
 			return nil, ErrUnsupported
 		}
 		u = sc.ChangeSchema(u, *cfg.schema)
 	}
-	client, err := v.(*driver).Open(ctx, u)
+	client, err := drv.Open(ctx, u)
 	if err != nil {
 		return nil, err
 	}
 	if client.URL == nil {
-		client.URL = v.(*driver).parser.ParseURL(u)
+		client.URL = drv.parser.ParseURL(u)
+	}
+	if client.openTx == nil && drv.txOpener != nil {
+		client.openTx = drv.txOpener
 	}
 	return client, nil
 }
@@ -332,15 +321,34 @@ func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error)) Opener
 }
 
 type (
-	// Tx wraps sql.Tx with an optional closer.
+	// Tx wraps sql.Tx with optional custom Commit and Rollback functions.
 	Tx struct {
 		*sql.Tx
-		Close func() error
+		CommitFn   func() error // override default commit behavior
+		RollbackFn func() error // override default rollback behavior
 	}
 	// TxOpener opens a transaction with optional closer.
 	TxOpener func(context.Context, *sql.DB, *sql.TxOptions) (*Tx, error)
 )
 
+// Commit the transaction.
+func (tx *Tx) Commit() error {
+	fn := tx.CommitFn
+	if fn == nil {
+		fn = tx.Tx.Commit
+	}
+	return fn()
+}
+
+// Rollback the transaction.
+func (tx *Tx) Rollback() error {
+	fn := tx.RollbackFn
+	if fn == nil {
+		fn = tx.Tx.Rollback
+	}
+	return fn()
+}
+
 // RegisterTxOpener allows registering a custom transaction opener with an optional close function.
 func RegisterTxOpener(open TxOpener) RegisterOption {
 	return func(opts *registerOptions) {
@@ -383,7 +391,7 @@ func Register(name string, opener Opener, opts ...RegisterOption) {
 			return c, err
 		})
 	}
-	drv := &driver{opener, name, opt.parser, opt.txOpener}
+	drv := &driver{Opener: opener, name: name, parser: opt.parser, txOpener: 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 93edd8ea..85a2b79f 100644
--- a/sql/sqlite/driver.go
+++ b/sql/sqlite/driver.go
@@ -48,7 +48,7 @@ func init() {
 	sqlclient.Register(
 		DriverName,
 		sqlclient.DriverOpener(Open),
-		sqlclient.RegisterTxOpener(openTx),
+		sqlclient.RegisterTxOpener(OpenTx),
 		sqlclient.RegisterCodec(MarshalHCL, EvalHCL),
 		sqlclient.RegisterFlavours("sqlite"),
 		sqlclient.RegisterURLParser(sqlclient.URLParserFunc(func(u *url.URL) *sqlclient.URL {
@@ -73,10 +73,10 @@ func Open(db schema.ExecQuerier) (migrate.Driver, error) {
 		return nil, fmt.Errorf("sqlite: query version pragma: %w", err)
 	}
 	if err := sqlx.ScanOne(rows, &c.version); err != nil {
-		return nil, fmt.Errorf("sqlite: scan version and foreign_keys pragma: %w", err)
+		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 foreign_keys pragma: %w", err)
+		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)
@@ -165,7 +165,13 @@ 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) {
+type violation struct {
+	tbl, ref   string
+	row, index int
+}
+
+// OpenTx opens a transaction. If foreign keys are enabled, it disables them, checks for constraint violations before doing so, disabling
+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)
@@ -181,19 +187,89 @@ func openTx(ctx context.Context, db *sql.DB, opts *sql.TxOptions) (*sqlclient.Tx
 	if err != nil {
 		return nil, err
 	}
+	before, err := violations(ctx, tx)
+	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)
+		CommitFn: func() error {
+			after, err := violations(ctx, tx)
+			if err != nil {
+				if err2 := tx.Rollback(); err2 != nil {
+					err = fmt.Errorf("%v: %w", err2, err)
+				}
+				return enableFK(ctx, db, on.Bool, err)
+			}
+			if vs := violationsDiff(before, after); len(vs) > 0 {
+				err := fmt.Errorf("sql/sqlite: foreign key mismatch: %+v", vs)
+				if err2 := tx.Rollback(); err2 != nil {
+					err = fmt.Errorf("%v: %w", err2, err)
 				}
+				return enableFK(ctx, db, on.Bool, err)
 			}
-			return nil
+			return enableFK(ctx, db, on.Bool, tx.Commit())
+		},
+		RollbackFn: func() error {
+			return enableFK(ctx, db, on.Bool, tx.Rollback())
 		},
 	}, nil
 }
 
+func enableFK(ctx context.Context, db *sql.DB, do bool, err error) error {
+	if do {
+		// Re-enable foreign key checks if they were enabled before.
+		if _, err2 := db.ExecContext(ctx, "PRAGMA foreign_keys = on"); err2 != nil {
+			err2 = fmt.Errorf("sql/sqlite: set 'foreign_keys = on': %w", err2)
+			if err != nil {
+				return fmt.Errorf("%v: %w", err2, err)
+			}
+			return err2
+		}
+	}
+	return err
+}
+
+func violations(ctx context.Context, tx *sql.Tx) ([]violation, error) {
+	rows, err := tx.QueryContext(ctx, "PRAGMA foreign_key_check")
+	if err != nil {
+		return nil, fmt.Errorf("sql/sqlite: querying 'foreign_key_check' pragma: %w", err)
+	}
+	defer rows.Close()
+	var vs []violation
+	for rows.Next() {
+		var v violation
+		if err := rows.Scan(&v.tbl, &v.row, &v.ref, &v.index); err != nil {
+			return nil, fmt.Errorf("sql/sqlite: querying 'foreign_key_check' pragma: scanning rows: %w", err)
+		}
+		vs = append(vs, v)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("sql/sqlite: querying 'foreign_key_check' pragma: scanning rows: %w", err)
+	}
+	return vs, nil
+}
+
+// equalViolations compares the foreign key violations before starting a transaction with the ones afterwards.
+// It returns violations found in v2 that are not in v1.
+func violationsDiff(v1, v2 []violation) (vs []violation) {
+	for _, v := range v2 {
+		if !contains(v1, v) {
+			vs = append(vs, v)
+		}
+	}
+	return vs
+}
+
+func contains(hs []violation, n violation) bool {
+	for _, v := range hs {
+		if v.row == n.row && v.ref == n.ref && v.index == n.index && v.tbl == n.tbl {
+			return true
+		}
+	}
+	return false
+}
+
 // 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


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate_test.go b/cmd/atlas/internal/cmdapi/migrate_test.go
index 10f5ad46..495b2dca 100644
--- a/cmd/atlas/internal/cmdapi/migrate_test.go
+++ b/cmd/atlas/internal/cmdapi/migrate_test.go
@@ -395,8 +395,38 @@ func TestMigrate_ApplyTxMode(t *testing.T) {
 				"--url", fmt.Sprintf("sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test.db")),
 				"--tx-mode", mode,
 			)
+			require.NoError(t, err)
 			require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM `friendships`").Scan(&n))
 			require.Equal(t, 2, n)
+
+			// For transactions check that the foreign keys are checked before the transaction is committed.
+			if mode != "none" {
+				// Apply the first 2 migrations for the faulty one.
+				s, err = runCmd(
+					Root, "migrate", "apply",
+					"--dir", "file://testdata/sqlitetx_2",
+					"--url", fmt.Sprintf("sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test_2.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_2.db")))
+				require.NoError(t, err)
+				require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM `friendships`").Scan(&n))
+				require.Equal(t, 2, n)
+
+				// Apply the rest, expect it to fail due to constraint error.
+				s, err = runCmd(
+					Root, "migrate", "apply",
+					"--dir", "file://testdata/sqlitetx_2",
+					"--url", fmt.Sprintf("sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test_2.db")),
+					"--tx-mode", mode,
+				)
+				require.ErrorContains(t, err, "sql/sqlite: foreign key mismatch: [{tbl:friendships ref:users row:3 index:1}]")
+				require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM `friendships`").Scan(&n))
+				require.Equal(t, 2, n) // was rolled back
+			}
 		})
 	}
 }
diff --git a/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/20220925092817_initial.sql b/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/20220925092817_initial.sql
new file mode 100644
index 00000000..ffd24610
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/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_2/20220925094021_second.sql b/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/20220925094021_second.sql
new file mode 100644
index 00000000..b2c69b70
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/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_2/20220925094437_third.sql b/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/20220925094437_third.sql
new file mode 100644
index 00000000..52368b94
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/20220925094437_third.sql
@@ -0,0 +1,14 @@
+-- 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`;
+-- insert faulty data
+INSERT INTO `friendships` (`user_id`, `friend_id`) VALUES (3,2);
+-- enable back the enforcement of foreign-keys constraints
+PRAGMA foreign_keys = on;
diff --git a/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/atlas.sum b/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/atlas.sum
new file mode 100644
index 00000000..4ea7a64f
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/sqlitetx_2/atlas.sum
@@ -0,0 +1,4 @@
+h1:eH+7c2mVWbrhky00/3SKYklyHLyz+QzDk1UZJ9+ZJsg=
+20220925092817_initial.sql h1:ZGeLdeqNUMXqJm+hPkhBrhzbtUzSBH8yVTsnSnJo/qU=
+20220925094021_second.sql h1:vcoquz3yk+TlTPiQgW5hHpS/abIvySCM/bzgwYTDoqY=
+20220925094437_third.sql h1:58glD96PVBa0fSu8x/3Gbwbf7N6WjAcVl4jh/XcNXoM=
diff --git a/sql/sqlclient/client_test.go b/sql/sqlclient/client_test.go
index 37a4b293..297d0a71 100644
--- a/sql/sqlclient/client_test.go
+++ b/sql/sqlclient/client_test.go
@@ -88,6 +88,7 @@ func TestClient_Tx(t *testing.T) {
 	mock.ExpectExec(stmt).WillReturnResult(sqlmock.NewResult(0, 1))
 	mock.ExpectRollback()
 
+	var cC, rC bool
 	sqlclient.Register(
 		"tx",
 		sqlclient.OpenerFunc(func(context.Context, *url.URL) (*sqlclient.Client, error) {
@@ -99,7 +100,17 @@ func TestClient_Tx(t *testing.T) {
 		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
+			return &sqlclient.Tx{
+				Tx: tx,
+				CommitFn: func() error {
+					cC = true
+					return tx.Commit()
+				},
+				RollbackFn: func() error {
+					rC = true
+					return tx.Rollback()
+				},
+			}, nil
 		}),
 	)
 
@@ -112,7 +123,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)
+	require.True(t, cC)
 
 	// Rollback works as well.
 	tx, err = c.Tx(context.Background(), nil)
@@ -120,6 +131,7 @@ func TestClient_Tx(t *testing.T) {
 	_, err = tx.ExecContext(context.Background(), stmt)
 	require.NoError(t, err)
 	require.NoError(t, tx.Rollback())
+	require.True(t, rC)
 
 	require.NoError(t, mock.ExpectationsWereMet())
 }