[pgx,postgres] Drop also drops custom PostgreSQL types

golang-migrate · Mar 11, 2024 · 5ed7ffa · 5ed7ffa
1 parent a5dc5d1
commit 5ed7ffa
Show file tree

Hide file tree

Showing 7 changed files with 161 additions and 21 deletions.
diff --git a/database/pgx/examples/migrations/1710098289_custom_type.down.sql b/database/pgx/examples/migrations/1710098289_custom_type.down.sql
@@ -0,0 +1,2 @@
+DROP TABLE IF EXISTS "custom_users";
+DROP TYPE IF EXISTS "custom_user_type";
diff --git a/database/pgx/examples/migrations/1710098289_custom_type.up.sql b/database/pgx/examples/migrations/1710098289_custom_type.up.sql
@@ -0,0 +1,8 @@
+CREATE TYPE "custom_user_type" AS ENUM('foo', 'bar', 'qux');
+
+CREATE TABLE "custom_users" (
+  "user_id"   integer unique,
+  "name"      text,
+  "email"     text,
+  "user_type" custom_user_type
+);
diff --git a/database/pgx/pgx.go b/database/pgx/pgx.go
@@ -535,13 +535,57 @@ func (p *Postgres) Drop() (err error) {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
 
-	if len(tableNames) > 0 {
-		// delete one by one ...
-		for _, t := range tableNames {
-			query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
-			if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
-				return &database.Error{OrigErr: err, Query: []byte(query)}
-			}
+	for _, t := range tableNames {
+		query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
+		if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
+			return &database.Error{OrigErr: err, Query: []byte(query)}
+		}
+	}
+
+	// select all types
+	query = `
+SELECT    t.typname as type
+FROM      pg_type t
+LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
+WHERE     (t.typrelid = 0 OR (SELECT c.relkind = 'c' FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid))
+AND       NOT EXISTS(SELECT 1 FROM pg_catalog.pg_type el WHERE el.oid = t.typelem AND el.typarray = t.oid)
+AND       n.nspname = current_schema()
+`
+	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 = multierror.Append(err, errClose)
+		}
+	}()
+
+	// delete one table after another
+	typeNames := make([]string, 0)
+	for types.Next() {
+		var typeName string
+		if err := types.Scan(&typeName); err != nil {
+			return err
+		}
+
+		// do not drop lock table
+		if typeName == p.config.LockTable && p.config.LockStrategy == LockStrategyTable {
+			continue
+		}
+
+		if len(typeName) > 0 {
+			typeNames = append(typeNames, typeName)
+		}
+	}
+	if err := types.Err(); err != nil {
+		return &database.Error{OrigErr: err, Query: []byte(query)}
+	}
+
+	for _, t := range typeNames {
+		query = `DROP TYPE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
+		if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
+			return &database.Error{OrigErr: err, Query: []byte(query)}
 		}
 	}
 

diff --git a/database/pgx/v5/pgx.go b/database/pgx/v5/pgx.go
@@ -420,13 +420,51 @@ func (p *Postgres) Drop() (err error) {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
 
-	if len(tableNames) > 0 {
-		// delete one by one ...
-		for _, t := range tableNames {
-			query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
-			if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
-				return &database.Error{OrigErr: err, Query: []byte(query)}
-			}
+	for _, t := range tableNames {
+		query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
+		if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
+			return &database.Error{OrigErr: err, Query: []byte(query)}
+		}
+	}
+
+	// select all custom types
+	query = `
+SELECT    t.typname as type
+FROM      pg_type t
+LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
+WHERE     (t.typrelid = 0 OR (SELECT c.relkind = 'c' FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid))
+AND       NOT EXISTS(SELECT 1 FROM pg_catalog.pg_type el WHERE el.oid = t.typelem AND el.typarray = t.oid)
+AND       n.nspname = current_schema()
+`
+	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 = multierror.Append(err, errClose)
+		}
+	}()
+
+	// delete one table after another
+	typeNames := make([]string, 0)
+	for types.Next() {
+		var tableName string
+		if err := types.Scan(&tableName); err != nil {
+			return err
+		}
+		if len(tableName) > 0 {
+			typeNames = append(typeNames, tableName)
+		}
+	}
+	if err := types.Err(); err != nil {
+		return &database.Error{OrigErr: err, Query: []byte(query)}
+	}
+
+	for _, t := range typeNames {
+		query = `DROP TYPE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
+		if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
+			return &database.Error{OrigErr: err, Query: []byte(query)}
 		}
 	}
 

diff --git a/database/postgres/examples/migrations/1710098289_custom_type.down.sql b/database/postgres/examples/migrations/1710098289_custom_type.down.sql
@@ -0,0 +1,2 @@
+DROP TABLE IF EXISTS "custom_users";
+DROP TYPE IF EXISTS "custom_user_type";
diff --git a/database/postgres/examples/migrations/1710098289_custom_type.up.sql b/database/postgres/examples/migrations/1710098289_custom_type.up.sql
@@ -0,0 +1,8 @@
+CREATE TYPE "custom_user_type" AS ENUM('foo', 'bar', 'qux');
+
+CREATE TABLE "custom_users" (
+  "user_id"   integer unique,
+  "name"      text,
+  "email"     text,
+  "user_type" custom_user_type
+);
diff --git a/database/postgres/postgres.go b/database/postgres/postgres.go
@@ -436,13 +436,51 @@ func (p *Postgres) Drop() (err error) {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
 
-	if len(tableNames) > 0 {
-		// delete one by one ...
-		for _, t := range tableNames {
-			query = `DROP TABLE IF EXISTS ` + pq.QuoteIdentifier(t) + ` CASCADE`
-			if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
-				return &database.Error{OrigErr: err, Query: []byte(query)}
-			}
+	for _, t := range tableNames {
+		query = `DROP TABLE IF EXISTS ` + pq.QuoteIdentifier(t) + ` CASCADE`
+		if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
+			return &database.Error{OrigErr: err, Query: []byte(query)}
+		}
+	}
+
+	// select all custom types
+	query = `
+SELECT    t.typname as type
+FROM      pg_type t
+LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
+WHERE     (t.typrelid = 0 OR (SELECT c.relkind = 'c' FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid))
+AND       NOT EXISTS(SELECT 1 FROM pg_catalog.pg_type el WHERE el.oid = t.typelem AND el.typarray = t.oid)
+AND       n.nspname = current_schema()
+`
+	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 = multierror.Append(err, errClose)
+		}
+	}()
+
+	// delete one type after another
+	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)}
+	}
+
+	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)}
 		}
 	}