Skip to content

Commit 5ed7ffa

Browse files
toqueteostoqueteos
committed
[pgx,postgres] Drop also drops custom PostgreSQL types
1 parent a5dc5d1 commit 5ed7ffa

File tree

7 files changed

+161
-21
lines changed

7 files changed

+161
-21
lines changed

‎database/pgx/examples/migrations/1710098289_custom_type.down.sql

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
DROP TABLE IF EXISTS "custom_users";
2+
DROP TYPE IF EXISTS "custom_user_type";

‎database/pgx/examples/migrations/1710098289_custom_type.up.sql

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
CREATE TYPE "custom_user_type" AS ENUM('foo', 'bar', 'qux');
2+
3+
CREATE TABLE "custom_users" (
4+
"user_id" integer unique,
5+
"name" text,
6+
"email" text,
7+
"user_type" custom_user_type
8+
);

‎database/pgx/pgx.go

+51-7
Original file line numberDiff line numberDiff line change
@@ -535,13 +535,57 @@ func (p *Postgres) Drop() (err error) {
535535
return &database.Error{OrigErr: err, Query: []byte(query)}
536536
}
537537

538-
if len(tableNames) > 0 {
539-
// delete one by one ...
540-
for _, t := range tableNames {
541-
query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
542-
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
543-
return &database.Error{OrigErr: err, Query: []byte(query)}
544-
}
538+
for _, t := range tableNames {
539+
query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
540+
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
541+
return &database.Error{OrigErr: err, Query: []byte(query)}
542+
}
543+
}
544+
545+
// select all types
546+
query = `
547+
SELECT t.typname as type
548+
FROM pg_type t
549+
LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
550+
WHERE (t.typrelid = 0 OR (SELECT c.relkind = 'c' FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid))
551+
AND NOT EXISTS(SELECT 1 FROM pg_catalog.pg_type el WHERE el.oid = t.typelem AND el.typarray = t.oid)
552+
AND n.nspname = current_schema()
553+
`
554+
types, err := p.conn.QueryContext(context.Background(), query)
555+
if err != nil {
556+
return &database.Error{OrigErr: err, Query: []byte(query)}
557+
}
558+
defer func() {
559+
if errClose := types.Close(); errClose != nil {
560+
err = multierror.Append(err, errClose)
561+
}
562+
}()
563+
564+
// delete one table after another
565+
typeNames := make([]string, 0)
566+
for types.Next() {
567+
var typeName string
568+
if err := types.Scan(&typeName); err != nil {
569+
return err
570+
}
571+
572+
// do not drop lock table
573+
if typeName == p.config.LockTable && p.config.LockStrategy == LockStrategyTable {
574+
continue
575+
}
576+
577+
if len(typeName) > 0 {
578+
typeNames = append(typeNames, typeName)
579+
}
580+
}
581+
if err := types.Err(); err != nil {
582+
return &database.Error{OrigErr: err, Query: []byte(query)}
583+
}
584+
585+
for _, t := range typeNames {
586+
query = `DROP TYPE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
587+
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
588+
return &database.Error{OrigErr: err, Query: []byte(query)}
545589
}
546590
}
547591

‎database/pgx/v5/pgx.go

+45-7
Original file line numberDiff line numberDiff line change
@@ -420,13 +420,51 @@ func (p *Postgres) Drop() (err error) {
420420
return &database.Error{OrigErr: err, Query: []byte(query)}
421421
}
422422

423-
if len(tableNames) > 0 {
424-
// delete one by one ...
425-
for _, t := range tableNames {
426-
query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
427-
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
428-
return &database.Error{OrigErr: err, Query: []byte(query)}
429-
}
423+
for _, t := range tableNames {
424+
query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
425+
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
426+
return &database.Error{OrigErr: err, Query: []byte(query)}
427+
}
428+
}
429+
430+
// select all custom types
431+
query = `
432+
SELECT t.typname as type
433+
FROM pg_type t
434+
LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
435+
WHERE (t.typrelid = 0 OR (SELECT c.relkind = 'c' FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid))
436+
AND NOT EXISTS(SELECT 1 FROM pg_catalog.pg_type el WHERE el.oid = t.typelem AND el.typarray = t.oid)
437+
AND n.nspname = current_schema()
438+
`
439+
types, err := p.conn.QueryContext(context.Background(), query)
440+
if err != nil {
441+
return &database.Error{OrigErr: err, Query: []byte(query)}
442+
}
443+
defer func() {
444+
if errClose := types.Close(); errClose != nil {
445+
err = multierror.Append(err, errClose)
446+
}
447+
}()
448+
449+
// delete one table after another
450+
typeNames := make([]string, 0)
451+
for types.Next() {
452+
var tableName string
453+
if err := types.Scan(&tableName); err != nil {
454+
return err
455+
}
456+
if len(tableName) > 0 {
457+
typeNames = append(typeNames, tableName)
458+
}
459+
}
460+
if err := types.Err(); err != nil {
461+
return &database.Error{OrigErr: err, Query: []byte(query)}
462+
}
463+
464+
for _, t := range typeNames {
465+
query = `DROP TYPE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE`
466+
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
467+
return &database.Error{OrigErr: err, Query: []byte(query)}
430468
}
431469
}
432470

‎database/postgres/examples/migrations/1710098289_custom_type.down.sql

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
DROP TABLE IF EXISTS "custom_users";
2+
DROP TYPE IF EXISTS "custom_user_type";

‎database/postgres/examples/migrations/1710098289_custom_type.up.sql

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
CREATE TYPE "custom_user_type" AS ENUM('foo', 'bar', 'qux');
2+
3+
CREATE TABLE "custom_users" (
4+
"user_id" integer unique,
5+
"name" text,
6+
"email" text,
7+
"user_type" custom_user_type
8+
);

‎database/postgres/postgres.go

+45-7
Original file line numberDiff line numberDiff line change
@@ -436,13 +436,51 @@ func (p *Postgres) Drop() (err error) {
436436
return &database.Error{OrigErr: err, Query: []byte(query)}
437437
}
438438

439-
if len(tableNames) > 0 {
440-
// delete one by one ...
441-
for _, t := range tableNames {
442-
query = `DROP TABLE IF EXISTS ` + pq.QuoteIdentifier(t) + ` CASCADE`
443-
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
444-
return &database.Error{OrigErr: err, Query: []byte(query)}
445-
}
439+
for _, t := range tableNames {
440+
query = `DROP TABLE IF EXISTS ` + pq.QuoteIdentifier(t) + ` CASCADE`
441+
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
442+
return &database.Error{OrigErr: err, Query: []byte(query)}
443+
}
444+
}
445+
446+
// select all custom types
447+
query = `
448+
SELECT t.typname as type
449+
FROM pg_type t
450+
LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
451+
WHERE (t.typrelid = 0 OR (SELECT c.relkind = 'c' FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid))
452+
AND NOT EXISTS(SELECT 1 FROM pg_catalog.pg_type el WHERE el.oid = t.typelem AND el.typarray = t.oid)
453+
AND n.nspname = current_schema()
454+
`
455+
types, err := p.conn.QueryContext(context.Background(), query)
456+
if err != nil {
457+
return &database.Error{OrigErr: err, Query: []byte(query)}
458+
}
459+
defer func() {
460+
if errClose := types.Close(); errClose != nil {
461+
err = multierror.Append(err, errClose)
462+
}
463+
}()
464+
465+
// delete one type after another
466+
typeNames := make([]string, 0)
467+
for types.Next() {
468+
var typeName string
469+
if err := types.Scan(&typeName); err != nil {
470+
return err
471+
}
472+
if len(typeName) > 0 {
473+
typeNames = append(typeNames, typeName)
474+
}
475+
}
476+
if err := types.Err(); err != nil {
477+
return &database.Error{OrigErr: err, Query: []byte(query)}
478+
}
479+
480+
for _, t := range typeNames {
481+
query = `DROP TYPE IF EXISTS ` + pq.QuoteIdentifier(t) + ` CASCADE`
482+
if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
483+
return &database.Error{OrigErr: err, Query: []byte(query)}
446484
}
447485
}
448486

0 commit comments

Comments
 (0)