From 5ed7ffabd77d1de65258c1b5cb0ce8f72bb00658 Mon Sep 17 00:00:00 2001 From: Carlos Cobo <699969+toqueteos@users.noreply.github.com> Date: Sun, 10 Mar 2024 19:31:22 +0000 Subject: [PATCH] [pgx,postgres] Drop also drops custom PostgreSQL types --- .../1710098289_custom_type.down.sql | 2 + .../migrations/1710098289_custom_type.up.sql | 8 +++ database/pgx/pgx.go | 58 ++++++++++++++++--- database/pgx/v5/pgx.go | 52 ++++++++++++++--- .../1710098289_custom_type.down.sql | 2 + .../migrations/1710098289_custom_type.up.sql | 8 +++ database/postgres/postgres.go | 52 ++++++++++++++--- 7 files changed, 161 insertions(+), 21 deletions(-) create mode 100644 database/pgx/examples/migrations/1710098289_custom_type.down.sql create mode 100644 database/pgx/examples/migrations/1710098289_custom_type.up.sql create mode 100644 database/postgres/examples/migrations/1710098289_custom_type.down.sql create mode 100644 database/postgres/examples/migrations/1710098289_custom_type.up.sql diff --git a/database/pgx/examples/migrations/1710098289_custom_type.down.sql b/database/pgx/examples/migrations/1710098289_custom_type.down.sql new file mode 100644 index 000000000..28d8db2c8 --- /dev/null +++ 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 new file mode 100644 index 000000000..d87f973b4 --- /dev/null +++ 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 index 7e42d29c9..faeb6355a 100644 --- 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 index 1b5a6ea7a..51f7eaacf 100644 --- 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 new file mode 100644 index 000000000..28d8db2c8 --- /dev/null +++ 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 new file mode 100644 index 000000000..d87f973b4 --- /dev/null +++ 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 index 9e6d6277f..2fc758b0b 100644 --- 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)} } }