PgTidy/pkg/lint/rules_migration.go

package lint

import (
	"fmt"
	"strings"

	pg_query "github.com/pganalyze/pg_query_go/v6"

	"git.warky.dev/wdevs/pgtidy/pkg/diagnostics"
	"git.warky.dev/wdevs/pgtidy/pkg/pgast"
)

// MIG001 — CREATE INDEX without CONCURRENT.
// A non-concurrent index build holds ShareLock for its entire duration,
// blocking all writes on the table. In production migrations use
// CREATE INDEX CONCURRENTLY instead.
type ruleMIG001 struct{}

func (ruleMIG001) ID() string                     { return "MIG001" }
func (ruleMIG001) Severity() diagnostics.Severity { return diagnostics.SeverityWarning }

func (ruleMIG001) Check(stmts []*pg_query.RawStmt, src string) []diagnostics.Diagnostic {
	var out []diagnostics.Diagnostic
	for _, raw := range stmts {
		idx, ok := raw.Stmt.GetNode().(*pg_query.Node_IndexStmt)
		if !ok {
			continue
		}
		s := idx.IndexStmt
		if s.Concurrent || s.Isconstraint {
			continue
		}
		line, col := pgast.LocationToLineCol(src, pgast.FirstTokenOffset(src, int(raw.StmtLocation)))
		d := diagnostics.Diagnostic{
			RuleID:   "MIG001",
			Severity: diagnostics.SeverityWarning,
			Message:  fmt.Sprintf("CREATE INDEX on %q without CONCURRENT blocks writes; use CREATE INDEX CONCURRENTLY", relName(s.Relation)),
			Line:     line,
			Col:      col,
		}
		if fix := mig001Fix(src, int(raw.StmtLocation), int(raw.StmtLen)); fix != nil {
			d.Fix = fix
		}
		out = append(out, d)
	}
	return out
}

// MIG002 — ALTER TABLE ADD COLUMN with NOT NULL and no DEFAULT.
// Pre-PG11 this rewrites the whole table. Even on PG11+ the column may require
// a costly table scan to validate the NOT NULL constraint when no DEFAULT is
// supplied. Use a nullable column first, backfill, then add the constraint.
type ruleMIG002 struct{}

func (ruleMIG002) ID() string                     { return "MIG002" }
func (ruleMIG002) Severity() diagnostics.Severity { return diagnostics.SeverityWarning }

func (ruleMIG002) Check(stmts []*pg_query.RawStmt, src string) []diagnostics.Diagnostic {
	var out []diagnostics.Diagnostic
	for _, raw := range stmts {
		alt, ok := raw.Stmt.GetNode().(*pg_query.Node_AlterTableStmt)
		if !ok {
			continue
		}
		tbl := alt.AlterTableStmt
		for _, cmdNode := range tbl.Cmds {
			cmd, ok := cmdNode.GetNode().(*pg_query.Node_AlterTableCmd)
			if !ok {
				continue
			}
			ac := cmd.AlterTableCmd
			if ac.Subtype != pg_query.AlterTableType_AT_AddColumn {
				continue
			}
			col, ok := ac.Def.GetNode().(*pg_query.Node_ColumnDef)
			if !ok {
				continue
			}
			if hasNotNullNoDefault(col.ColumnDef.Constraints) {
				line, col2 := pgast.LocationToLineCol(src, pgast.FirstTokenOffset(src, int(raw.StmtLocation)))
				out = append(out, diagnostics.Diagnostic{
					RuleID:   "MIG002",
					Severity: diagnostics.SeverityWarning,
					Message:  fmt.Sprintf("ALTER TABLE %q ADD COLUMN with NOT NULL and no DEFAULT may rewrite the table; add nullable, backfill, then constrain", relName(tbl.Relation)),
					Line:     line,
					Col:      col2,
				})
			}
		}
	}
	return out
}

// MIG003 — ALTER TABLE ADD CONSTRAINT (FK or CHECK) without NOT VALID.
// Without NOT VALID, PostgreSQL validates all existing rows immediately,
// holding locks that block concurrent writes. Use NOT VALID + a separate
// VALIDATE CONSTRAINT to spread the lock window.
type ruleMIG003 struct{}

func (ruleMIG003) ID() string                     { return "MIG003" }
func (ruleMIG003) Severity() diagnostics.Severity { return diagnostics.SeverityWarning }

func (ruleMIG003) Check(stmts []*pg_query.RawStmt, src string) []diagnostics.Diagnostic {
	var out []diagnostics.Diagnostic
	for _, raw := range stmts {
		alt, ok := raw.Stmt.GetNode().(*pg_query.Node_AlterTableStmt)
		if !ok {
			continue
		}
		tbl := alt.AlterTableStmt
		for _, cmdNode := range tbl.Cmds {
			cmd, ok := cmdNode.GetNode().(*pg_query.Node_AlterTableCmd)
			if !ok {
				continue
			}
			if cmd.AlterTableCmd.Subtype != pg_query.AlterTableType_AT_AddConstraint {
				continue
			}
			con, ok := cmd.AlterTableCmd.Def.GetNode().(*pg_query.Node_Constraint)
			if !ok {
				continue
			}
			c := con.Constraint
			if c.Contype != pg_query.ConstrType_CONSTR_FOREIGN && c.Contype != pg_query.ConstrType_CONSTR_CHECK {
				continue
			}
			if c.SkipValidation {
				continue // NOT VALID is present
			}
			kind := "FOREIGN KEY"
			if c.Contype == pg_query.ConstrType_CONSTR_CHECK {
				kind = "CHECK"
			}
			line, col := pgast.LocationToLineCol(src, pgast.FirstTokenOffset(src, int(raw.StmtLocation)))
			d := diagnostics.Diagnostic{
				RuleID:   "MIG003",
				Severity: diagnostics.SeverityWarning,
				Message:  fmt.Sprintf("ALTER TABLE %q ADD %s CONSTRAINT without NOT VALID validates all rows immediately; use NOT VALID + VALIDATE CONSTRAINT", relName(tbl.Relation), kind),
				Line:     line,
				Col:      col,
			}
			if fix := mig003Fix(src, int(raw.StmtLocation), int(raw.StmtLen)); fix != nil {
				d.Fix = fix
			}
			out = append(out, d)
		}
	}
	return out
}

// hasNotNullNoDefault returns true when the constraint list has CONSTR_NOTNULL
// but no CONSTR_DEFAULT.
func hasNotNullNoDefault(constraints []*pg_query.Node) bool {
	hasNN := false
	hasDef := false
	for _, cn := range constraints {
		c, ok := cn.GetNode().(*pg_query.Node_Constraint)
		if !ok {
			continue
		}
		switch c.Constraint.Contype {
		case pg_query.ConstrType_CONSTR_NOTNULL:
			hasNN = true
		case pg_query.ConstrType_CONSTR_DEFAULT:
			hasDef = true
		}
	}
	return hasNN && !hasDef
}

// relName returns "schema.rel" or "rel" from a RangeVar.
func relName(rv *pg_query.RangeVar) string {
	if rv == nil {
		return "?"
	}
	if rv.Schemaname != "" {
		return rv.Schemaname + "." + rv.Relname
	}
	return rv.Relname
}

// stmtText returns the text of a statement given its start offset and length.
// When stmtLen is 0 (last statement in file) it extends to EOF.
func stmtText(src string, stmtOffset, stmtLen int) string {
	end := stmtOffset + stmtLen
	if stmtLen == 0 || end > len(src) {
		end = len(src)
	}
	return src[stmtOffset:end]
}

// mig001Fix builds the TextFix for MIG001: insert CONCURRENTLY after INDEX.
func mig001Fix(src string, stmtOffset, stmtLen int) *diagnostics.TextFix {
	text := stmtText(src, stmtOffset, stmtLen)
	upper := strings.ToUpper(text)
	pos := strings.Index(upper, "INDEX")
	if pos < 0 {
		return nil
	}
	insertAt := stmtOffset + pos + len("INDEX")
	return &diagnostics.TextFix{
		Offset: insertAt,
		End:    insertAt,
		New:    " CONCURRENTLY",
		Title:  "Add CONCURRENTLY",
	}
}

// mig003Fix builds the TextFix for MIG003: insert NOT VALID before the trailing semicolon.
// pg_query's StmtLen excludes the ";", which sits at src[stmtOffset+stmtLen].
func mig003Fix(src string, stmtOffset, stmtLen int) *diagnostics.TextFix {
	end := stmtOffset + stmtLen
	semiAt := -1
	if stmtLen > 0 && end < len(src) && src[end] == ';' {
		semiAt = end
	} else {
		// Fallback for stmtLen=0 (last statement without terminator) or edge cases.
		if stmtLen == 0 {
			end = len(src)
		}
		idx := strings.LastIndex(src[stmtOffset:end], ";")
		if idx >= 0 {
			semiAt = stmtOffset + idx
		}
	}
	if semiAt < 0 {
		return nil
	}
	// Insert " NOT VALID" just before the ";", after any trailing whitespace.
	insertAt := semiAt
	for insertAt > stmtOffset && (src[insertAt-1] == ' ' || src[insertAt-1] == '\t' || src[insertAt-1] == '\n' || src[insertAt-1] == '\r') {
		insertAt--
	}
	return &diagnostics.TextFix{
		Offset: insertAt,
		End:    insertAt,
		New:    " NOT VALID",
		Title:  "Add NOT VALID",
	}
}