Repository: ariga/atlas
Pull request #974 (2022-07-21 20:30:06)

sql/postgres: delete dangling enum types


Patch:
diff --git a/sql/postgres/convert.go b/sql/postgres/convert.go
index e38b627f..fd590f36 100644
--- a/sql/postgres/convert.go
+++ b/sql/postgres/convert.go
@@ -181,6 +181,86 @@ func ParseType(typ string) (schema.Type, error) {
 	return t, nil
 }
 
+func columnType(c *columnDesc) (schema.Type, error) {
+	var typ schema.Type
+	switch t := c.typ; strings.ToLower(t) {
+	case TypeBigInt, TypeInt8, TypeInt, TypeInteger, TypeInt4, TypeSmallInt, TypeInt2, TypeInt64:
+		typ = &schema.IntegerType{T: t}
+	case TypeBit, TypeBitVar:
+		typ = &BitType{T: t, Len: c.size}
+	case TypeBool, TypeBoolean:
+		typ = &schema.BoolType{T: t}
+	case TypeBytea:
+		typ = &schema.BinaryType{T: t}
+	case TypeCharacter, TypeChar, TypeCharVar, TypeVarChar, TypeText:
+		// A `character` column without length specifier is equivalent to `character(1)`,
+		// but `varchar` without length accepts strings of any size (same as `text`).
+		typ = &schema.StringType{T: t, Size: int(c.size)}
+	case TypeCIDR, TypeInet, TypeMACAddr, TypeMACAddr8:
+		typ = &NetworkType{T: t}
+	case TypeCircle, TypeLine, TypeLseg, TypeBox, TypePath, TypePolygon, TypePoint, TypeGeometry:
+		typ = &schema.SpatialType{T: t}
+	case TypeDate:
+		typ = &schema.TimeType{T: t}
+	case TypeTime, TypeTimeWOTZ, TypeTimeTZ, TypeTimeWTZ, TypeTimestamp,
+		TypeTimestampTZ, TypeTimestampWTZ, TypeTimestampWOTZ:
+		p := defaultTimePrecision
+		if c.timePrecision != nil {
+			p = int(*c.timePrecision)
+		}
+		typ = &schema.TimeType{T: t, Precision: &p}
+	case TypeInterval:
+		p := defaultTimePrecision
+		if c.timePrecision != nil {
+			p = int(*c.timePrecision)
+		}
+		typ = &IntervalType{T: t, Precision: &p}
+		if c.interval != "" {
+			f, ok := intervalField(c.interval)
+			if !ok {
+				return &schema.UnsupportedType{T: c.interval}, nil
+			}
+			typ.(*IntervalType).F = f
+		}
+	case TypeReal, TypeDouble, TypeFloat, TypeFloat4, TypeFloat8:
+		typ = &schema.FloatType{T: t, Precision: int(c.precision)}
+	case TypeJSON, TypeJSONB:
+		typ = &schema.JSONType{T: t}
+	case TypeMoney:
+		typ = &CurrencyType{T: t}
+	case TypeDecimal, TypeNumeric:
+		typ = &schema.DecimalType{T: t, Precision: int(c.precision), Scale: int(c.scale)}
+	case TypeSmallSerial, TypeSerial, TypeBigSerial, TypeSerial2, TypeSerial4, TypeSerial8:
+		typ = &SerialType{T: t, Precision: int(c.precision)}
+	case TypeUUID:
+		typ = &UUIDType{T: t}
+	case TypeXML:
+		typ = &XMLType{T: t}
+	case TypeArray:
+		// Ignore multi-dimensions or size constraints
+		// as they are ignored by the database.
+		typ = &ArrayType{T: c.fmtype}
+		if t, ok := arrayType(c.fmtype); ok {
+			tt, err := ParseType(t)
+			if err != nil {
+				return nil, err
+			}
+			typ.(*ArrayType).Type = tt
+		}
+	case TypeUserDefined:
+		typ = &UserDefinedType{T: c.fmtype}
+		// The `typtype` column is set to 'e' for enum types, and the
+		// values are filled in batch after the rows above is closed.
+		// https://postgresql.org/docs/current/catalog-pg-type.html
+		if c.typtype == "e" {
+			typ = newEnumType(c.fmtype, c.typid)
+		}
+	default:
+		typ = &schema.UnsupportedType{T: t}
+	}
+	return typ, nil
+}
+
 // reArray parses array declaration. See: https://postgresql.org/docs/current/arrays.html.
 var reArray = regexp.MustCompile(`(?i)(.+?)(( +ARRAY( *\[[ \d]*] *)*)+|( *\[[ \d]*] *)+)$`)
 
diff --git a/sql/postgres/diff.go b/sql/postgres/diff.go
index d2fec87f..2d8418fd 100644
--- a/sql/postgres/diff.go
+++ b/sql/postgres/diff.go
@@ -254,56 +254,6 @@ func (d *diff) typeChanged(from, to *schema.Column) (bool, error) {
 	return changed, nil
 }
 
-// Normalize implements the sqlx.Normalizer interface.
-func (d *diff) Normalize(from, to *schema.Table) error {
-	d.normalize(from)
-	d.normalize(to)
-	return nil
-}
-
-func (d *diff) normalize(table *schema.Table) {
-	for _, c := range table.Columns {
-		switch t := c.Type.Type.(type) {
-		case nil:
-		case *schema.TimeType:
-			// "timestamp" and "timestamptz" are accepted as
-			// abbreviations for timestamp with(out) time zone.
-			switch t.T {
-			case "timestamp with time zone":
-				t.T = "timestamptz"
-			case "timestamp without time zone":
-				t.T = "timestamp"
-			}
-		case *schema.FloatType:
-			// The same numeric precision is used in all platform.
-			// See: https://www.postgresql.org/docs/current/datatype-numeric.html
-			switch {
-			case t.T == "float" && t.Precision < 25:
-				// float(1) to float(24) are selected as "real" type.
-				t.T = "real"
-				fallthrough
-			case t.T == "real":
-				t.Precision = 24
-			case t.T == "float" && t.Precision >= 25:
-				// float(25) to float(53) are selected as "double precision" type.
-				t.T = "double precision"
-				fallthrough
-			case t.T == "double precision":
-				t.Precision = 53
-			}
-		case *schema.StringType:
-			switch t.T {
-			case "character", "char":
-				// Character without length specifier
-				// is equivalent to character(1).
-				t.Size = 1
-			}
-		case *enumType:
-			c.Type.Type = &schema.EnumType{T: t.T, Values: t.Values}
-		}
-	}
-}
-
 // valuesEqual reports if the DEFAULT values x and y
 // equal according to the database engine.
 func (d *diff) valuesEqual(x, y string) (bool, error) {
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 215758a0..6cb4bb48 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -228,86 +228,6 @@ func (i *inspect) addColumn(s *schema.Schema, rows *sql.Rows) (err error) {
 	return nil
 }
 
-func columnType(c *columnDesc) (schema.Type, error) {
-	var typ schema.Type
-	switch t := c.typ; strings.ToLower(t) {
-	case TypeBigInt, TypeInt8, TypeInt, TypeInteger, TypeInt4, TypeSmallInt, TypeInt2, TypeInt64:
-		typ = &schema.IntegerType{T: t}
-	case TypeBit, TypeBitVar:
-		typ = &BitType{T: t, Len: c.size}
-	case TypeBool, TypeBoolean:
-		typ = &schema.BoolType{T: t}
-	case TypeBytea:
-		typ = &schema.BinaryType{T: t}
-	case TypeCharacter, TypeChar, TypeCharVar, TypeVarChar, TypeText:
-		// A `character` column without length specifier is equivalent to `character(1)`,
-		// but `varchar` without length accepts strings of any size (same as `text`).
-		typ = &schema.StringType{T: t, Size: int(c.size)}
-	case TypeCIDR, TypeInet, TypeMACAddr, TypeMACAddr8:
-		typ = &NetworkType{T: t}
-	case TypeCircle, TypeLine, TypeLseg, TypeBox, TypePath, TypePolygon, TypePoint, TypeGeometry:
-		typ = &schema.SpatialType{T: t}
-	case TypeDate:
-		typ = &schema.TimeType{T: t}
-	case TypeTime, TypeTimeWOTZ, TypeTimeTZ, TypeTimeWTZ, TypeTimestamp,
-		TypeTimestampTZ, TypeTimestampWTZ, TypeTimestampWOTZ:
-		p := defaultTimePrecision
-		if c.timePrecision != nil {
-			p = int(*c.timePrecision)
-		}
-		typ = &schema.TimeType{T: t, Precision: &p}
-	case TypeInterval:
-		p := defaultTimePrecision
-		if c.timePrecision != nil {
-			p = int(*c.timePrecision)
-		}
-		typ = &IntervalType{T: t, Precision: &p}
-		if c.interval != "" {
-			f, ok := intervalField(c.interval)
-			if !ok {
-				return &schema.UnsupportedType{T: c.interval}, nil
-			}
-			typ.(*IntervalType).F = f
-		}
-	case TypeReal, TypeDouble, TypeFloat, TypeFloat4, TypeFloat8:
-		typ = &schema.FloatType{T: t, Precision: int(c.precision)}
-	case TypeJSON, TypeJSONB:
-		typ = &schema.JSONType{T: t}
-	case TypeMoney:
-		typ = &CurrencyType{T: t}
-	case TypeDecimal, TypeNumeric:
-		typ = &schema.DecimalType{T: t, Precision: int(c.precision), Scale: int(c.scale)}
-	case TypeSmallSerial, TypeSerial, TypeBigSerial, TypeSerial2, TypeSerial4, TypeSerial8:
-		typ = &SerialType{T: t, Precision: int(c.precision)}
-	case TypeUUID:
-		typ = &UUIDType{T: t}
-	case TypeXML:
-		typ = &XMLType{T: t}
-	case TypeArray:
-		// Ignore multi-dimensions or size constraints
-		// as they are ignored by the database.
-		typ = &ArrayType{T: c.fmtype}
-		if t, ok := arrayType(c.fmtype); ok {
-			tt, err := ParseType(t)
-			if err != nil {
-				return nil, err
-			}
-			typ.(*ArrayType).Type = tt
-		}
-	case TypeUserDefined:
-		typ = &UserDefinedType{T: c.fmtype}
-		// The `typtype` column is set to 'e' for enum types, and the
-		// values are filled in batch after the rows above is closed.
-		// https://postgresql.org/docs/current/catalog-pg-type.html
-		if c.typtype == "e" {
-			typ = newEnumType(c.fmtype, c.typid)
-		}
-	default:
-		typ = &schema.UnsupportedType{T: t}
-	}
-	return typ, nil
-}
-
 // enumValues fills enum columns with their values from the database.
 func (i *inspect) enumValues(ctx context.Context, s *schema.Schema) error {
 	var (
@@ -936,19 +856,6 @@ func newEnumType(t string, id int64) *enumType {
 	return e
 }
 
-func enumIdent(s *schema.Schema, e *schema.EnumType) string {
-	switch {
-	// Enum schema has higher precedence.
-	case e.Schema != nil:
-		return fmt.Sprintf("%q.%q", e.Schema.Name, e.T)
-	// Fallback to table schema if exists.
-	case s != nil:
-		return fmt.Sprintf("%q.%q", s.Name, e.T)
-	default:
-		return strconv.Quote(e.T)
-	}
-}
-
 const (
 	// Query to list runtime parameters.
 	paramsQuery = `SELECT setting FROM pg_settings WHERE name IN ('lc_collate', 'lc_ctype', 'server_version_num', 'crdb_version') ORDER BY name DESC`
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 74241637..60ae2961 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -126,7 +126,7 @@ func (s *state) topLevel(changes []schema.Change) []schema.Change {
 // addTable builds and executes the query for creating a table in a schema.
 func (s *state) addTable(ctx context.Context, add *schema.AddTable) error {
 	// Create enum types before using them in the `CREATE TABLE` statement.
-	if err := s.addTypes(ctx, add.T, add.T.Columns...); err != nil {
+	if err := s.mayAddEnums(ctx, add.T, add.T.Columns...); err != nil {
 		return err
 	}
 	var (
@@ -255,7 +255,7 @@ func (s *state) modifyTable(ctx context.Context, modify *schema.ModifyTable) err
 				F: change.To,
 			})
 		case *schema.AddColumn:
-			if err := s.addTypes(ctx, modify.T, change.C); err != nil {
+			if err := s.mayAddEnums(ctx, modify.T, change.C); err != nil {
 				return err
 			}
 			if c := (schema.Comment{}); sqlx.Has(change.C.Attrs, &c) {
@@ -278,7 +278,7 @@ func (s *state) modifyTable(ctx context.Context, modify *schema.ModifyTable) err
 			from, ok1 := change.From.Type.Type.(*schema.EnumType)
 			to, ok2 := change.To.Type.Type.(*schema.EnumType)
 			switch {
-			// Enum was changed.
+			// Enum was changed (underlying values).
 			case ok1 && ok2 && enumIdent(modify.T.Schema, from) == enumIdent(modify.T.Schema, to):
 				if err := s.alterEnum(modify.T, from, to); err != nil {
 					return err
@@ -288,10 +288,10 @@ func (s *state) modifyTable(ctx context.Context, modify *schema.ModifyTable) err
 				if k == schema.ChangeType {
 					continue
 				}
-			// Column type was changed to enum, or enum type was changed.
+			// Enum was added or changed.
 			case !ok1 && ok2 ||
 				ok1 && ok2 && enumIdent(modify.T.Schema, from) != enumIdent(modify.T.Schema, to):
-				if err := s.addTypes(ctx, modify.T, change.To); err != nil {
+				if err := s.mayAddEnums(ctx, modify.T, change.To); err != nil {
 					return err
 				}
 			}
@@ -349,9 +349,22 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 					To:     change.From,
 					Change: change.Change & ^schema.ChangeGenerated,
 				})
+				toE, toHas := change.From.Type.Type.(*schema.EnumType)
+				fromE, fromHas := change.From.Type.Type.(*schema.EnumType)
+				// In case the enum was dropped or replaced with a different one.
+				if fromHas && !toHas || fromHas && toHas && enumIdent(t.Schema, fromE) != enumIdent(t.Schema, toE) {
+					if err := mayDropEnum(alter, t.Schema, fromE); err != nil {
+						return err
+					}
+				}
 			case *schema.DropColumn:
 				b.P("DROP COLUMN").Ident(change.C.Name)
 				reverse = append(reverse, &schema.AddColumn{C: change.C})
+				if e, ok := change.C.Type.Type.(*schema.EnumType); ok {
+					if err := mayDropEnum(alter, t.Schema, e); err != nil {
+						return err
+					}
+				}
 			case *schema.AddIndex:
 				b.P("ADD CONSTRAINT").Ident(change.I.Name).P("UNIQUE")
 				s.indexParts(b, change.I.Parts)
@@ -650,7 +663,7 @@ func (s *state) dropIndexes(t *schema.Table, indexes ...*schema.Index) {
 	}
 }
 
-func (s *state) addTypes(ctx context.Context, t *schema.Table, columns ...*schema.Column) error {
+func (s *state) mayAddEnums(ctx context.Context, t *schema.Table, columns ...*schema.Column) error {
 	for _, c := range columns {
 		e, ok := c.Type.Type.(*schema.EnumType)
 		if !ok {
@@ -674,16 +687,11 @@ func (s *state) addTypes(ctx context.Context, t *schema.Table, columns ...*schem
 			continue
 		}
 		s.created[name] = e
-		b := Build("CREATE TYPE").P(name, "AS ENUM")
-		b.Wrap(func(b *sqlx.Builder) {
-			b.MapComma(e.Values, func(i int, b *sqlx.Builder) {
-				b.WriteString(quote(e.Values[i]))
-			})
-		})
+		create, drop := createDropEnum(t.Schema, e)
 		s.append(&migrate.Change{
-			Cmd:     b.String(),
+			Cmd:     create,
+			Reverse: drop,
 			Comment: fmt.Sprintf("create enum type %q", e.T),
-			Reverse: Build("DROP TYPE").P(name).String(),
 		})
 	}
 	return nil
@@ -719,7 +727,7 @@ func (s *state) enumExists(ctx context.Context, ns *schema.Schema, e *schema.Enu
 	if e.Schema != nil {
 		ns = e.Schema
 	}
-	rows, err := s.QueryContext(ctx, `SELECT * FROM pg_type t join pg_namespace n on t.typnamespace = n.oid WHERE t.typname = $1 AND n.nspname = $2 AND t.typtype = 'e'`, e.T, ns.Name)
+	rows, err := s.QueryContext(ctx, `SELECT * FROM pg_type t JOIN pg_namespace n on t.typnamespace = n.oid WHERE t.typname = $1 AND n.nspname = $2 AND t.typtype = 'e'`, e.T, ns.Name)
 	if err != nil {
 		return false, fmt.Errorf("check enum existence: %w", err)
 	}
@@ -727,6 +735,34 @@ func (s *state) enumExists(ctx context.Context, ns *schema.Schema, e *schema.Enu
 	return rows.Next(), rows.Err()
 }
 
+// mayDropEnum drops dangling enum types form the schema.
+func mayDropEnum(alter *alterChange, s *schema.Schema, e *schema.EnumType) error {
+	schemas := []*schema.Schema{s}
+	// In case there is a realm attached, traverse the entire tree.
+	if s.Realm != nil && len(s.Realm.Schemas) > 0 {
+		schemas = s.Realm.Schemas
+	}
+	for _, ns := range schemas {
+		for _, t := range ns.Tables {
+			for _, c := range t.Columns {
+				e1, ok := c.Type.Type.(*schema.EnumType)
+				// Although we search in siblings schemas, use the
+				// table's one for building the enum identifier.
+				if ok && enumIdent(s, e1) == enumIdent(s, e) {
+					return nil
+				}
+			}
+		}
+	}
+	create, drop := createDropEnum(s, e)
+	alter.after = append(alter.after, &migrate.Change{
+		Cmd:     drop,
+		Reverse: create,
+		Comment: fmt.Sprintf("drop enum type %q", e.T),
+	})
+	return nil
+}
+
 func (s *state) addIndexes(t *schema.Table, indexes ...*schema.Index) {
 	for _, idx := range indexes {
 		b := Build("CREATE")
@@ -1046,3 +1082,29 @@ func quote(s string) string {
 	}
 	return "'" + strings.ReplaceAll(s, "'", "''") + "'"
 }
+
+func enumIdent(s *schema.Schema, e *schema.EnumType) string {
+	switch {
+	// Enum schema has higher precedence.
+	case e.Schema != nil:
+		return fmt.Sprintf("%q.%q", e.Schema.Name, e.T)
+	// Fallback to table schema if exists.
+	case s != nil:
+		return fmt.Sprintf("%q.%q", s.Name, e.T)
+	default:
+		return strconv.Quote(e.T)
+	}
+}
+
+func createDropEnum(s *schema.Schema, e *schema.EnumType) (string, string) {
+	name := enumIdent(s, e)
+	return Build("CREATE TYPE").
+			P(name, "AS ENUM").
+			Wrap(func(b *sqlx.Builder) {
+				b.MapComma(e.Values, func(i int, b *sqlx.Builder) {
+					b.WriteString(quote(e.Values[i]))
+				})
+			}).
+			String(),
+		Build("DROP TYPE").P(name).String()
+}


Test patch:
diff --git a/internal/integration/testdata/postgres/column-enum.txt b/internal/integration/testdata/postgres/column-enum.txt
index 613b7744..22039fc9 100644
--- a/internal/integration/testdata/postgres/column-enum.txt
+++ b/internal/integration/testdata/postgres/column-enum.txt
@@ -1,10 +1,18 @@
 apply 1.hcl
 cmpshow users 1.sql
 
-# Change enum default.
 apply 2.hcl
 cmpshow users 2.sql
 
+# Drop the enum.
+apply 3.hcl
+cmpshow users 3.sql
+
+# Add it back.
+apply 2.hcl
+cmpshow users 2.sql
+
+
 -- 1.hcl --
 schema "$db" {}
 
@@ -49,3 +57,20 @@ table "users" {
  Column |           Type            | Collation | Nullable |                Default
 --------+---------------------------+-----------+----------+---------------------------------------
  type   | script_column_enum.status |           | not null | 'inactive'::script_column_enum.status
+
+-- 3.hcl --
+schema "$db" {}
+
+table "users" {
+  schema = schema.$db
+  column "int" {
+    type = int
+    default = 1
+  }
+}
+
+-- 3.sql --
+         Table "script_column_enum.users"
+ Column |  Type   | Collation | Nullable | Default
+--------+---------+-----------+----------+---------
+ int    | integer |           | not null | 1
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index d70ad388..52cb0d5b 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -540,15 +540,18 @@ func TestPlanChanges(t *testing.T) {
 								To:     users.Columns[1],
 								Change: schema.ChangeType,
 							},
+							&schema.DropColumn{
+								C: schema.NewEnumColumn("dc", schema.EnumName("de"), schema.EnumValues("on")),
+							},
 						},
 					}
 				}(),
 			},
 			mock: func(m mock) {
-				m.ExpectQuery(sqltest.Escape("SELECT * FROM pg_type t join pg_namespace n on t.typnamespace = n.oid WHERE t.typname = $1 AND n.nspname = $2 AND t.typtype = 'e'")).
+				m.ExpectQuery(sqltest.Escape("SELECT * FROM pg_type t JOIN pg_namespace n on t.typnamespace = n.oid WHERE t.typname = $1 AND n.nspname = $2 AND t.typtype = 'e'")).
 					WithArgs("state", "public").
 					WillReturnRows(sqlmock.NewRows([]string{"name"}))
-				m.ExpectQuery(sqltest.Escape("SELECT * FROM pg_type t join pg_namespace n on t.typnamespace = n.oid WHERE t.typname = $1 AND n.nspname = $2 AND t.typtype = 'e'")).
+				m.ExpectQuery(sqltest.Escape("SELECT * FROM pg_type t JOIN pg_namespace n on t.typnamespace = n.oid WHERE t.typname = $1 AND n.nspname = $2 AND t.typtype = 'e'")).
 					WithArgs("status", "test").
 					WillReturnRows(sqlmock.NewRows([]string{"name"}))
 			},
@@ -558,7 +561,8 @@ func TestPlanChanges(t *testing.T) {
 				Changes: []*migrate.Change{
 					{Cmd: `CREATE TYPE "public"."state" AS ENUM ('on', 'off')`, Reverse: `DROP TYPE "public"."state"`},
 					{Cmd: `CREATE TYPE "test"."status" AS ENUM ('a', 'b')`, Reverse: `DROP TYPE "test"."status"`},
-					{Cmd: `ALTER TABLE "public"."users" ALTER COLUMN "state" TYPE "public"."state", ALTER COLUMN "status" TYPE "test"."status"`, Reverse: `ALTER TABLE "public"."users" ALTER COLUMN "status" TYPE text, ALTER COLUMN "state" TYPE text`},
+					{Cmd: `ALTER TABLE "public"."users" ALTER COLUMN "state" TYPE "public"."state", ALTER COLUMN "status" TYPE "test"."status", DROP COLUMN "dc"`, Reverse: `ALTER TABLE "public"."users" ADD COLUMN "dc" "public"."de" NOT NULL, ALTER COLUMN "status" TYPE text, ALTER COLUMN "state" TYPE text`},
+					{Cmd: `DROP TYPE "public"."de"`, Reverse: `CREATE TYPE "public"."de" AS ENUM ('on')`},
 				},
 			},
 		},