Repository: ariga/atlas
Pull request #1559 (2023-03-27 08:24:31)

sql/postgres/postgrescheck: add concurrent index commands analyzer
An analyzer that detects two cases:
1. Missing the `CONCURRENTLY` option in `CREATE INDEX`/`DROP INDEX` commands and report accordingly. Error (red CI) if was set this way. Note the code skips tables that were created in the same file.
2. Warning about missing `atlas:txmode none` directive in case the migration file contains concurrent commands (invalid operation).

Future improvements: As concurrent creation and deletion operations cannot run within a transaction, it is recommended that users remove unrelated statements (such as `ALTER TABLE`) from the file. This is because other commands can, and should, be executed within a transaction.


Patch:
diff --git a/cmd/atlas/internal/sqlparse/pgparse/pgparse.go b/cmd/atlas/internal/sqlparse/pgparse/pgparse.go
index 33498599..4183d1e3 100644
--- a/cmd/atlas/internal/sqlparse/pgparse/pgparse.go
+++ b/cmd/atlas/internal/sqlparse/pgparse/pgparse.go
@@ -112,7 +112,7 @@ func (p *Parser) FixChange(_ migrate.Driver, s string, changes schema.Changes) (
 		if slices.IndexFunc(add.Extra, func(c schema.Clause) bool {
 			_, ok := c.(*postgres.Concurrently)
 			return ok
-		}) == -1 {
+		}) == -1 && stmt.Concurrently {
 			add.Extra = append(add.Extra, &postgres.Concurrently{})
 		}
 	}
diff --git a/sql/postgres/postgrescheck/postgrescheck.go b/sql/postgres/postgrescheck/postgrescheck.go
index e4358ea1..c3dc4625 100644
--- a/sql/postgres/postgrescheck/postgrescheck.go
+++ b/sql/postgres/postgrescheck/postgrescheck.go
@@ -5,10 +5,15 @@
 package postgrescheck
 
 import (
+	"context"
+	"errors"
 	"fmt"
 
 	"ariga.io/atlas/schemahcl"
+	"ariga.io/atlas/sql/internal/sqlx"
+	"ariga.io/atlas/sql/migrate"
 	"ariga.io/atlas/sql/postgres"
+	"ariga.io/atlas/sql/schema"
 	"ariga.io/atlas/sql/sqlcheck"
 	"ariga.io/atlas/sql/sqlcheck/condrop"
 	"ariga.io/atlas/sql/sqlcheck/datadepend"
@@ -54,6 +59,136 @@ func init() {
 		if err != nil {
 			return nil, err
 		}
-		return []sqlcheck.Analyzer{ds, dd, cd, bc, nm}, nil
+		ci, err := NewConcurrentIndex(r)
+		if err != nil {
+			return nil, err
+		}
+		return []sqlcheck.Analyzer{ds, dd, cd, bc, nm, ci}, nil
 	})
 }
+
+type (
+	// ConcurrentIndex checks for concurrent index creation,
+	// dropping, and its transactional safety.
+	ConcurrentIndex struct {
+		ConcurrentOptions
+	}
+	// ConcurrentOptions for concurrent index creation.
+	ConcurrentOptions struct {
+		sqlcheck.Options
+		CheckCreate *bool `spec:"check_create"`
+		CheckDrop   *bool `spec:"check_drop"`
+		CheckTxMode *bool `spec:"check_txmode"`
+	}
+)
+
+// NewConcurrentIndex creates a new concurrent-index Analyzer with the given options.
+func NewConcurrentIndex(r *schemahcl.Resource) (*ConcurrentIndex, error) {
+	az := &ConcurrentIndex{}
+	az.CheckCreate = sqlx.P(true)
+	az.CheckDrop = sqlx.P(true)
+	az.CheckTxMode = sqlx.P(true)
+	if r, ok := r.Resource(az.Name()); ok {
+		if err := r.As(&az.Options); err != nil {
+			return nil, fmt.Errorf("sql/sqlcheck: parsing concurrent_index error option: %w", err)
+		}
+		if err := r.As(&az.ConcurrentOptions); err != nil {
+			return nil, fmt.Errorf("sql/sqlcheck: parsing concurrent_index check options: %w", err)
+		}
+	}
+	return az, nil
+}
+
+// Name of the analyzer. Implements the sqlcheck.NamedAnalyzer interface.
+func (*ConcurrentIndex) Name() string {
+	return "concurrent_index"
+}
+
+var (
+	// codeCreateNoCon is a PostgreSQL specific code for reporting
+	// indexes creation without the CONCURRENTLY clause.
+	codeCreateNoCon = sqlcheck.Code("PG101")
+	// codeDropNoCon is a PostgreSQL specific code for reporting
+	// indexes deletion without the CONCURRENTLY clause.
+	codeDropNoCon = sqlcheck.Code("PG102")
+	// codeNoTxNone is a PostgreSQL specific code for reporting indexes
+	// creation or deletion concurrently without the txmode directive set
+	// to none.
+	codeNoTxNone = sqlcheck.Code("PG103")
+)
+
+// Analyze implements sqlcheck.Analyzer.
+func (a *ConcurrentIndex) Analyze(_ context.Context, p *sqlcheck.Pass) error {
+	var (
+		notx  bool
+		notxC int
+		diags []sqlcheck.Diagnostic
+	)
+	// The txmode directive, currently defined in cmd/atlas,
+	// might be moved to sql/migrate in the future.
+	if l, ok := p.File.File.(*migrate.LocalFile); ok {
+		mode := l.Directive("txmode")
+		notx = len(mode) == 1 && mode[0] == "none"
+	}
+	for _, sc := range p.File.Changes {
+		for _, c := range sc.Changes {
+			m, ok := c.(*schema.ModifyTable)
+			// Skip modifications for tables that have been created in this file.
+			if !ok || p.File.TableSpan(m.T)&sqlcheck.SpanAdded == 1 {
+				continue
+			}
+			for i := range m.Changes {
+				switch mc := m.Changes[i].(type) {
+				case *schema.AddIndex:
+					switch hasC := sqlx.Has(mc.Extra, &postgres.Concurrently{}); {
+					case !sqlx.V(a.CheckCreate):
+					case hasC && !notx && sqlx.V(a.CheckTxMode):
+						notxC++
+					case !hasC:
+						diags = append(diags, sqlcheck.Diagnostic{
+							Pos:  sc.Stmt.Pos,
+							Code: codeCreateNoCon,
+							Text: fmt.Sprintf(
+								"Creating index %q non-concurrently causes write locks on the %q table",
+								mc.I.Name, m.T.Name,
+							),
+						})
+					}
+				case *schema.DropIndex:
+					switch hasC := sqlx.Has(mc.Extra, &postgres.Concurrently{}); {
+					case !sqlx.V(a.CheckDrop):
+					case hasC && !notx && sqlx.V(a.CheckTxMode):
+						notxC++
+					case !hasC:
+						diags = append(diags, sqlcheck.Diagnostic{
+							Pos:  sc.Stmt.Pos,
+							Code: codeDropNoCon,
+							Text: fmt.Sprintf(
+								"Dropping index %q non-concurrently causes write locks on the %q table",
+								mc.I.Name, m.T.Name,
+							),
+						})
+					}
+				}
+			}
+		}
+	}
+	if notxC > 0 {
+		diags = append([]sqlcheck.Diagnostic{{
+			Pos:  0,
+			Code: codeNoTxNone,
+			Text: "Indexes cannot be created or deleted within a transaction. Add the `atlas:txmode none` " +
+				"directive to the header to prevent this file from running in a transaction",
+		}}, diags...)
+	}
+	if len(diags) > 0 {
+		const reportText = "concurrent index violations detected"
+		p.Reporter.WriteReport(sqlcheck.Report{Text: reportText, Diagnostics: diags})
+		// Report an error only if it is configured this way and the
+		// diagnostics include non-concurrent creation or deletion.
+		if sqlx.V(a.Error) && (notxC == 0 || len(diags) > 1) {
+			return errors.New(reportText)
+		}
+	}
+	return nil
+}


Test patch:
diff --git a/cmd/atlas/internal/sqlparse/pgparse/pgparse_test.go b/cmd/atlas/internal/sqlparse/pgparse/pgparse_test.go
index 4ee7e7e4..a8eb9d09 100644
--- a/cmd/atlas/internal/sqlparse/pgparse/pgparse_test.go
+++ b/cmd/atlas/internal/sqlparse/pgparse/pgparse_test.go
@@ -89,6 +89,35 @@ func TestFixChange_RenameIndexes(t *testing.T) {
 func TestFixChange_CreateIndexCon(t *testing.T) {
 	var p pgparse.Parser
 	changes, err := p.FixChange(
+		nil,
+		"CREATE INDEX i1 ON t1 (c1)",
+		schema.Changes{
+			&schema.ModifyTable{
+				T: schema.NewTable("t1"),
+				Changes: schema.Changes{
+					&schema.AddIndex{I: schema.NewIndex("i1")},
+				},
+			},
+		},
+	)
+	require.NoError(t, err)
+	// No changes.
+	require.Equal(
+		t,
+		schema.Changes{
+			&schema.ModifyTable{
+				T: schema.NewTable("t1"),
+				Changes: schema.Changes{
+					&schema.AddIndex{
+						I: schema.NewIndex("i1"),
+					},
+				},
+			},
+		},
+		changes,
+	)
+
+	changes, err = p.FixChange(
 		nil,
 		"CREATE INDEX CONCURRENTLY i1 ON t1 (c1)",
 		schema.Changes{
@@ -119,6 +148,7 @@ func TestFixChange_CreateIndexCon(t *testing.T) {
 		},
 		changes,
 	)
+
 	changes, err = p.FixChange(
 		nil,
 		"CREATE INDEX CONCURRENTLY i1 ON t1 (c1)",
diff --git a/internal/integration/testdata/postgres/cli-migrate-lint-concurrent-idx.txt b/internal/integration/testdata/postgres/cli-migrate-lint-concurrent-idx.txt
new file mode 100644
index 00000000..2616ae06
--- /dev/null
+++ b/internal/integration/testdata/postgres/cli-migrate-lint-concurrent-idx.txt
@@ -0,0 +1,52 @@
+only postgres15
+
+# Ignore tables that were created in the same file.
+atlas migrate lint --dir file://migrations1 --dev-url URL --latest=1 > got.txt
+cmp got.txt empty.txt
+
+atlas migrate lint --dir file://migrations2 --dev-url URL --latest=1 > got.txt
+cmp got.txt expected2.txt
+
+atlas migrate lint --dir file://migrations3 --dev-url URL --latest=1 > got.txt
+cmp got.txt expected3.txt
+
+atlas migrate lint --dir file://migrations4 --dev-url URL --latest=1 > got.txt
+cmp got.txt expected4.txt
+
+-- empty.txt --
+-- migrations1/1.sql --
+CREATE TABLE t(c int);
+CREATE INDEX i ON t(c);
+
+-- migrations2/1.sql --
+CREATE TABLE t(c int);
+CREATE INDEX i1 ON t(c);
+-- migrations2/2.sql --
+DROP INDEX i1;
+CREATE INDEX i2 ON t(c);
+-- expected2.txt --
+2.sql: concurrent index violations detected:
+
+	L1: Dropping index "i1" non-concurrently causes write locks on the "t" table
+	L2: Creating index "i2" non-concurrently causes write locks on the "t" table
+
+-- migrations3/1.sql --
+CREATE TABLE t(c int);
+-- migrations3/2.sql --
+CREATE INDEX CONCURRENTLY i ON t(c);
+-- expected3.txt --
+2.sql: concurrent index violations detected:
+
+	L1: Indexes cannot be created or deleted within a transaction. Add the `atlas:txmode none` directive to the header to prevent this file from running in a transaction
+
+-- migrations4/1.sql --
+CREATE TABLE t(c int);
+-- migrations4/2.sql --
+CREATE INDEX CONCURRENTLY i2 ON t(c);
+CREATE INDEX i1 ON t(c);
+-- expected4.txt --
+2.sql: concurrent index violations detected:
+
+	L1: Indexes cannot be created or deleted within a transaction. Add the `atlas:txmode none` directive to the header to prevent this file from running in a transaction
+	L2: Creating index "i1" non-concurrently causes write locks on the "t" table
+
diff --git a/sql/postgres/postgrescheck/postgrescheck_test.go b/sql/postgres/postgrescheck/postgrescheck_test.go
index f8beb66d..28e2efa6 100644
--- a/sql/postgres/postgrescheck/postgrescheck_test.go
+++ b/sql/postgres/postgrescheck/postgrescheck_test.go
@@ -8,8 +8,10 @@ import (
 	"context"
 	"testing"
 
+	"ariga.io/atlas/schemahcl"
 	"ariga.io/atlas/sql/migrate"
 	"ariga.io/atlas/sql/postgres"
+	"ariga.io/atlas/sql/postgres/postgrescheck"
 	_ "ariga.io/atlas/sql/postgres/postgrescheck"
 	"ariga.io/atlas/sql/schema"
 	"ariga.io/atlas/sql/sqlcheck"
@@ -17,6 +19,144 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestConcurrentIndex(t *testing.T) {
+	var cfg struct {
+		schemahcl.DefaultExtension
+	}
+	// language=hcl
+	err := schemahcl.New().EvalBytes([]byte(`
+concurrent_index {
+  error = true
+}
+`), &cfg, nil)
+	require.NoError(t, err)
+	require.NoError(t, err)
+	az, err := postgrescheck.NewConcurrentIndex(cfg.Remain())
+	require.NoError(t, err)
+
+	t.Run("MissingConcurrent", func(t *testing.T) {
+		var report *sqlcheck.Report
+		err := az.Analyze(context.Background(), &sqlcheck.Pass{
+			File: &sqlcheck.File{
+				File: migrate.NewLocalFile("1.sql", []byte("CREATE INDEX i1 ON t(c);")),
+				Changes: []*sqlcheck.Change{
+					{
+						Changes: schema.Changes{
+							&schema.ModifyTable{
+								T: schema.NewTable("Users").SetSchema(schema.New("public")),
+								Changes: schema.Changes{
+									&schema.AddIndex{
+										I: schema.NewIndex("i1"),
+									},
+								},
+							},
+						},
+						Stmt: &migrate.Stmt{
+							Pos:  0,
+							Text: "CREATE INDEX i1 ON t(c)",
+						},
+					},
+				},
+			},
+			Reporter: sqlcheck.ReportWriterFunc(func(r sqlcheck.Report) {
+				report = &r
+			}),
+		})
+		require.EqualError(t, err, "concurrent index violations detected")
+		require.Len(t, report.Diagnostics, 1)
+		require.Equal(t, report.Diagnostics[0].Text, `Creating index "i1" non-concurrently causes write locks on the "Users" table`)
+	})
+
+	t.Run("MissingTxMode", func(t *testing.T) {
+		var report *sqlcheck.Report
+		err := az.Analyze(context.Background(), &sqlcheck.Pass{
+			File: &sqlcheck.File{
+				File: migrate.NewLocalFile("1.sql", []byte("CREATE INDEX CONCURRENTLY i1 ON t(c);")),
+				Changes: []*sqlcheck.Change{
+					{
+						Changes: schema.Changes{
+							&schema.ModifyTable{
+								T: schema.NewTable("Users").SetSchema(schema.New("public")),
+								Changes: schema.Changes{
+									&schema.AddIndex{
+										I: schema.NewIndex("i1"),
+										Extra: []schema.Clause{
+											&postgres.Concurrently{},
+										},
+									},
+								},
+							},
+						},
+						Stmt: &migrate.Stmt{
+							Pos:  0,
+							Text: "CREATE INDEX CONCURRENTLY i1 ON t(c)",
+						},
+					},
+				},
+			},
+			Reporter: sqlcheck.ReportWriterFunc(func(r sqlcheck.Report) {
+				report = &r
+			}),
+		})
+		require.NoError(t, err)
+		require.Len(t, report.Diagnostics, 1)
+		require.Equal(t, report.Diagnostics[0].Text, "Indexes cannot be created or deleted within a transaction. Add the `atlas:txmode none` directive to the header to prevent this file from running in a transaction")
+	})
+
+	t.Run("MixedReport", func(t *testing.T) {
+		var report *sqlcheck.Report
+		err := az.Analyze(context.Background(), &sqlcheck.Pass{
+			File: &sqlcheck.File{
+				File: migrate.NewLocalFile("1.sql", []byte("CREATE INDEX CONCURRENTLY i1 ON t(c);DROP INDEX i2;")),
+				Changes: []*sqlcheck.Change{
+					{
+						Changes: schema.Changes{
+							&schema.ModifyTable{
+								T: schema.NewTable("Users").SetSchema(schema.New("public")),
+								Changes: schema.Changes{
+									&schema.AddIndex{
+										I: schema.NewIndex("i1"),
+										Extra: []schema.Clause{
+											&postgres.Concurrently{},
+										},
+									},
+								},
+							},
+						},
+						Stmt: &migrate.Stmt{
+							Pos:  0,
+							Text: "CREATE INDEX CONCURRENTLY i1 ON t(c)",
+						},
+					},
+					{
+						Changes: schema.Changes{
+							&schema.ModifyTable{
+								T: schema.NewTable("Users").SetSchema(schema.New("public")),
+								Changes: schema.Changes{
+									&schema.DropIndex{
+										I: schema.NewIndex("i2"),
+									},
+								},
+							},
+						},
+						Stmt: &migrate.Stmt{
+							Pos:  0,
+							Text: "DROP INDEX i2",
+						},
+					},
+				},
+			},
+			Reporter: sqlcheck.ReportWriterFunc(func(r sqlcheck.Report) {
+				report = &r
+			}),
+		})
+		require.EqualError(t, err, "concurrent index violations detected")
+		require.Len(t, report.Diagnostics, 2)
+		require.Equal(t, report.Diagnostics[0].Text, "Indexes cannot be created or deleted within a transaction. Add the `atlas:txmode none` directive to the header to prevent this file from running in a transaction")
+		require.Equal(t, report.Diagnostics[1].Text, `Dropping index "i2" non-concurrently causes write locks on the "Users" table`)
+	})
+}
+
 func TestDataDepend_MightFail(t *testing.T) {
 	var (
 		report *sqlcheck.Report