postgres: Drop() also drops custom types

database/postgres/postgres.go

@@ -444,6 +444,57 @@ func (p *Postgres) Drop() (err error) {
 		}
 	}
 
+	// select all custom types (enums, domains, and standalone composite types) in the current schema.
+	// Excludes the row type Postgres automatically creates for every table (already gone via DROP TABLE
+	// above, or never a real user type to begin with), the array type Postgres automatically creates
+	// alongside every type, and any type an extension or the system installed (an internal or
+	// extension pg_depend entry) -- those can't be dropped without dropping the extension itself,
+	// which is outside Drop()'s remit, and a range type's auto-generated multirange has an internal
+	// dependency on its range type that trips a plain DROP regardless of statement order.
+	query = `SELECT t.typname FROM pg_catalog.pg_type t
+		LEFT JOIN pg_catalog.pg_class c ON c.oid = t.typrelid
+		WHERE t.typnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = current_schema())
+		AND (t.typrelid = 0 OR c.relkind = 'c')
+		AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_type el WHERE el.typarray = t.oid)
+		AND NOT EXISTS (
+			SELECT 1 FROM pg_catalog.pg_depend d
+			WHERE d.classid = 'pg_catalog.pg_type'::regclass AND d.objid = t.oid
+			AND d.deptype IN ('e', 'i')
+		)`
+	types, err := p.conn.QueryContext(context.Background(), query)
+	if err != nil {
+		return &database.Error{OrigErr: err, Query: []byte(query)}
+	}
+	defer func() {
+		if errClose := types.Close(); errClose != nil {
+			err = errors.Join(err, errClose)
+		}
+	}()
+
+	typeNames := make([]string, 0)
+	for types.Next() {
+		var typeName string
+		if err := types.Scan(&typeName); err != nil {
+			return err
+		}
+		if len(typeName) > 0 {
+			typeNames = append(typeNames, typeName)
+		}
+	}
+	if err := types.Err(); err != nil {
+		return &database.Error{OrigErr: err, Query: []byte(query)}
+	}
+
+	if len(typeNames) > 0 {
+		// delete one by one ...
+		for _, t := range typeNames {
+			query = `DROP TYPE IF EXISTS ` + pq.QuoteIdentifier(t) + ` CASCADE`
+			if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
+				return &database.Error{OrigErr: err, Query: []byte(query)}
+			}
+		}
+	}
+
 	return nil
 }
 

database/postgres/postgres_test.go

@@ -99,6 +99,8 @@ func Test(t *testing.T) {
 	t.Run("testPostgresLock", testPostgresLock)
 	t.Run("testWithInstanceConcurrent", testWithInstanceConcurrent)
 	t.Run("testWithConnection", testWithConnection)
+	t.Run("testDropWithCustomTypes", testDropWithCustomTypes)
+	t.Run("testDropWithExtensionAndRangeTypes", testDropWithExtensionAndRangeTypes)
 
 	t.Cleanup(func() {
 		for _, spec := range specs {
@@ -132,6 +134,130 @@ func test(t *testing.T) {
 	})
 }
 
+// testDropWithCustomTypes is a regression test for Drop() not removing custom
+// types (enums, domains, standalone composite types), which left a database
+// unable to re-run a migration that (re-)creates one of those types. See
+// https://github.com/golang-migrate/migrate/issues/626.
+func testDropWithCustomTypes(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
+		ip, port, err := c.FirstPort()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		addr := pgConnectionString(ip, port)
+		p := &Postgres{}
+		d, err := p.Open(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := d.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+
+		db, err := sql.Open("postgres", addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := db.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+
+		statements := []string{
+			`CREATE TYPE mood AS ENUM ('sad', 'ok', 'happy')`,
+			`CREATE TYPE point2d AS (x int, y int)`,
+			`CREATE DOMAIN posint AS int CHECK (VALUE > 0)`,
+			`CREATE TABLE widgets (id serial primary key, m mood, p point2d, n posint)`,
+		}
+		for _, s := range statements {
+			if _, err := db.Exec(s); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		if err := d.Drop(); err != nil {
+			t.Fatal(err)
+		}
+
+		// A truly clean database can re-create every type Drop() should have
+		// removed. Before the fix this failed with "type ... already exists".
+		for _, s := range statements[:3] {
+			if _, err := db.Exec(s); err != nil {
+				t.Fatalf("re-creating type after Drop() failed (Drop() left it behind): %v", err)
+			}
+		}
+	})
+}
+
+// testDropWithExtensionAndRangeTypes is a regression test for two ways the
+// type-selection query in Drop() can pick up a type it must not try to drop:
+// a type an extension installed (dropping it fails -- only dropping the
+// extension itself can remove it), and a range type, whose auto-generated
+// multirange type has an internal dependency on it that a plain DROP TYPE
+// (even with CASCADE) cannot satisfy regardless of statement order.
+func testDropWithExtensionAndRangeTypes(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
+		ip, port, err := c.FirstPort()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		addr := pgConnectionString(ip, port)
+		p := &Postgres{}
+		d, err := p.Open(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := d.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+
+		db, err := sql.Open("postgres", addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := db.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+
+		statements := []string{
+			`CREATE EXTENSION IF NOT EXISTS citext`,
+			`CREATE TYPE mood AS ENUM ('sad', 'ok', 'happy')`,
+			`CREATE TYPE floatrange AS RANGE (subtype = float8)`,
+		}
+		for _, s := range statements {
+			if _, err := db.Exec(s); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		if err := d.Drop(); err != nil {
+			t.Fatalf("Drop() failed on a database with an extension-owned type and a range type: %v", err)
+		}
+
+		var extant string
+		err = db.QueryRow(`SELECT typname FROM pg_type WHERE typname = 'citext'`).Scan(&extant)
+		if err != nil {
+			t.Fatalf("expected the extension-owned citext type to survive Drop(), got: %v", err)
+		}
+
+		if _, err := db.Exec(`CREATE TYPE mood AS ENUM ('sad', 'ok', 'happy')`); err != nil {
+			t.Fatalf("re-creating mood after Drop() failed: %v", err)
+		}
+		if _, err := db.Exec(`CREATE TYPE floatrange AS RANGE (subtype = float8)`); err != nil {
+			t.Fatalf("re-creating floatrange after Drop() failed: %v", err)
+		}
+	})
+}
+
 func testMigrate(t *testing.T) {
 	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
 		ip, port, err := c.FirstPort()