github.hscsec.cn/amacneil/dbmate@v1.4.1/pkg/dbmate/postgres.go (about) 1 package dbmate 2 3 import ( 4 "bytes" 5 "database/sql" 6 "fmt" 7 "net/url" 8 "strings" 9 10 "github.com/lib/pq" 11 ) 12 13 func init() { 14 RegisterDriver(PostgresDriver{}, "postgres") 15 RegisterDriver(PostgresDriver{}, "postgresql") 16 } 17 18 // PostgresDriver provides top level database functions 19 type PostgresDriver struct { 20 } 21 22 // Open creates a new database connection 23 func (drv PostgresDriver) Open(u *url.URL) (*sql.DB, error) { 24 return sql.Open("postgres", u.String()) 25 } 26 27 func (drv PostgresDriver) openPostgresDB(u *url.URL) (*sql.DB, error) { 28 // connect to postgres database 29 postgresURL := *u 30 postgresURL.Path = "postgres" 31 32 return drv.Open(&postgresURL) 33 } 34 35 // CreateDatabase creates the specified database 36 func (drv PostgresDriver) CreateDatabase(u *url.URL) error { 37 name := databaseName(u) 38 fmt.Printf("Creating: %s\n", name) 39 40 db, err := drv.openPostgresDB(u) 41 if err != nil { 42 return err 43 } 44 defer mustClose(db) 45 46 _, err = db.Exec(fmt.Sprintf("create database %s", 47 pq.QuoteIdentifier(name))) 48 49 return err 50 } 51 52 // DropDatabase drops the specified database (if it exists) 53 func (drv PostgresDriver) DropDatabase(u *url.URL) error { 54 name := databaseName(u) 55 fmt.Printf("Dropping: %s\n", name) 56 57 db, err := drv.openPostgresDB(u) 58 if err != nil { 59 return err 60 } 61 defer mustClose(db) 62 63 _, err = db.Exec(fmt.Sprintf("drop database if exists %s", 64 pq.QuoteIdentifier(name))) 65 66 return err 67 } 68 69 func postgresSchemaMigrationsDump(db *sql.DB) ([]byte, error) { 70 // load applied migrations 71 migrations, err := queryColumn(db, 72 "select quote_literal(version) from public.schema_migrations order by version asc") 73 if err != nil { 74 return nil, err 75 } 76 77 // build schema_migrations table data 78 var buf bytes.Buffer 79 buf.WriteString("\n--\n-- Dbmate schema migrations\n--\n\n") 80 81 if len(migrations) > 0 { 82 buf.WriteString("INSERT INTO public.schema_migrations (version) VALUES\n (" + 83 strings.Join(migrations, "),\n (") + 84 ");\n") 85 } 86 87 return buf.Bytes(), nil 88 } 89 90 // DumpSchema returns the current database schema 91 func (drv PostgresDriver) DumpSchema(u *url.URL, db *sql.DB) ([]byte, error) { 92 // load schema 93 schema, err := runCommand("pg_dump", "--format=plain", "--encoding=UTF8", 94 "--schema-only", "--no-privileges", "--no-owner", u.String()) 95 if err != nil { 96 return nil, err 97 } 98 99 migrations, err := postgresSchemaMigrationsDump(db) 100 if err != nil { 101 return nil, err 102 } 103 104 schema = append(schema, migrations...) 105 return trimLeadingSQLComments(schema) 106 } 107 108 // DatabaseExists determines whether the database exists 109 func (drv PostgresDriver) DatabaseExists(u *url.URL) (bool, error) { 110 name := databaseName(u) 111 112 db, err := drv.openPostgresDB(u) 113 if err != nil { 114 return false, err 115 } 116 defer mustClose(db) 117 118 exists := false 119 err = db.QueryRow("select true from pg_database where datname = $1", name). 120 Scan(&exists) 121 if err == sql.ErrNoRows { 122 return false, nil 123 } 124 125 return exists, err 126 } 127 128 // CreateMigrationsTable creates the schema_migrations table 129 func (drv PostgresDriver) CreateMigrationsTable(db *sql.DB) error { 130 _, err := db.Exec("create table if not exists public.schema_migrations " + 131 "(version varchar(255) primary key)") 132 133 return err 134 } 135 136 // SelectMigrations returns a list of applied migrations 137 // with an optional limit (in descending order) 138 func (drv PostgresDriver) SelectMigrations(db *sql.DB, limit int) (map[string]bool, error) { 139 query := "select version from public.schema_migrations order by version desc" 140 if limit >= 0 { 141 query = fmt.Sprintf("%s limit %d", query, limit) 142 } 143 rows, err := db.Query(query) 144 if err != nil { 145 return nil, err 146 } 147 148 defer mustClose(rows) 149 150 migrations := map[string]bool{} 151 for rows.Next() { 152 var version string 153 if err := rows.Scan(&version); err != nil { 154 return nil, err 155 } 156 157 migrations[version] = true 158 } 159 160 return migrations, nil 161 } 162 163 // InsertMigration adds a new migration record 164 func (drv PostgresDriver) InsertMigration(db Transaction, version string) error { 165 _, err := db.Exec("insert into public.schema_migrations (version) values ($1)", version) 166 167 return err 168 } 169 170 // DeleteMigration removes a migration record 171 func (drv PostgresDriver) DeleteMigration(db Transaction, version string) error { 172 _, err := db.Exec("delete from public.schema_migrations where version = $1", version) 173 174 return err 175 } 176 177 // Ping verifies a connection to the database server. It does not verify whether the 178 // specified database exists. 179 func (drv PostgresDriver) Ping(u *url.URL) error { 180 db, err := drv.openPostgresDB(u) 181 if err != nil { 182 return err 183 } 184 defer mustClose(db) 185 186 return db.Ping() 187 }