github.com/kubecost/golang-migrate-duckdb/v4@v4.17.0-duckdb.1/database/pgx/pgx.go (about) 1 //go:build go1.9 2 // +build go1.9 3 4 package pgx 5 6 import ( 7 "context" 8 "database/sql" 9 "fmt" 10 "io" 11 nurl "net/url" 12 "regexp" 13 "strconv" 14 "strings" 15 "time" 16 17 "go.uber.org/atomic" 18 19 "github.com/golang-migrate/migrate/v4" 20 "github.com/golang-migrate/migrate/v4/database" 21 "github.com/golang-migrate/migrate/v4/database/multistmt" 22 "github.com/hashicorp/go-multierror" 23 "github.com/jackc/pgconn" 24 "github.com/jackc/pgerrcode" 25 _ "github.com/jackc/pgx/v4/stdlib" 26 ) 27 28 func init() { 29 db := Postgres{} 30 database.Register("pgx", &db) 31 database.Register("pgx4", &db) 32 } 33 34 var ( 35 multiStmtDelimiter = []byte(";") 36 37 DefaultMigrationsTable = "schema_migrations" 38 DefaultMultiStatementMaxSize = 10 * 1 << 20 // 10 MB 39 ) 40 41 var ( 42 ErrNilConfig = fmt.Errorf("no config") 43 ErrNoDatabaseName = fmt.Errorf("no database name") 44 ErrNoSchema = fmt.Errorf("no schema") 45 ErrDatabaseDirty = fmt.Errorf("database is dirty") 46 ) 47 48 type Config struct { 49 MigrationsTable string 50 DatabaseName string 51 SchemaName string 52 migrationsSchemaName string 53 migrationsTableName string 54 StatementTimeout time.Duration 55 MigrationsTableQuoted bool 56 MultiStatementEnabled bool 57 MultiStatementMaxSize int 58 } 59 60 type Postgres struct { 61 // Locking and unlocking need to use the same connection 62 conn *sql.Conn 63 db *sql.DB 64 isLocked atomic.Bool 65 66 // Open and WithInstance need to guarantee that config is never nil 67 config *Config 68 } 69 70 func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) { 71 if config == nil { 72 return nil, ErrNilConfig 73 } 74 75 if err := instance.Ping(); err != nil { 76 return nil, err 77 } 78 79 if config.DatabaseName == "" { 80 query := `SELECT CURRENT_DATABASE()` 81 var databaseName string 82 if err := instance.QueryRow(query).Scan(&databaseName); err != nil { 83 return nil, &database.Error{OrigErr: err, Query: []byte(query)} 84 } 85 86 if len(databaseName) == 0 { 87 return nil, ErrNoDatabaseName 88 } 89 90 config.DatabaseName = databaseName 91 } 92 93 if config.SchemaName == "" { 94 query := `SELECT CURRENT_SCHEMA()` 95 var schemaName string 96 if err := instance.QueryRow(query).Scan(&schemaName); err != nil { 97 return nil, &database.Error{OrigErr: err, Query: []byte(query)} 98 } 99 100 if len(schemaName) == 0 { 101 return nil, ErrNoSchema 102 } 103 104 config.SchemaName = schemaName 105 } 106 107 if len(config.MigrationsTable) == 0 { 108 config.MigrationsTable = DefaultMigrationsTable 109 } 110 111 config.migrationsSchemaName = config.SchemaName 112 config.migrationsTableName = config.MigrationsTable 113 if config.MigrationsTableQuoted { 114 re := regexp.MustCompile(`"(.*?)"`) 115 result := re.FindAllStringSubmatch(config.MigrationsTable, -1) 116 config.migrationsTableName = result[len(result)-1][1] 117 if len(result) == 2 { 118 config.migrationsSchemaName = result[0][1] 119 } else if len(result) > 2 { 120 return nil, fmt.Errorf("\"%s\" MigrationsTable contains too many dot characters", config.MigrationsTable) 121 } 122 } 123 124 conn, err := instance.Conn(context.Background()) 125 126 if err != nil { 127 return nil, err 128 } 129 130 px := &Postgres{ 131 conn: conn, 132 db: instance, 133 config: config, 134 } 135 136 if err := px.ensureVersionTable(); err != nil { 137 return nil, err 138 } 139 140 return px, nil 141 } 142 143 func (p *Postgres) Open(url string) (database.Driver, error) { 144 purl, err := nurl.Parse(url) 145 if err != nil { 146 return nil, err 147 } 148 149 // Driver is registered as pgx, but connection string must use postgres schema 150 // when making actual connection 151 // i.e. pgx://user:password@host:port/db => postgres://user:password@host:port/db 152 purl.Scheme = "postgres" 153 154 db, err := sql.Open("pgx/v4", migrate.FilterCustomQuery(purl).String()) 155 if err != nil { 156 return nil, err 157 } 158 159 migrationsTable := purl.Query().Get("x-migrations-table") 160 migrationsTableQuoted := false 161 if s := purl.Query().Get("x-migrations-table-quoted"); len(s) > 0 { 162 migrationsTableQuoted, err = strconv.ParseBool(s) 163 if err != nil { 164 return nil, fmt.Errorf("Unable to parse option x-migrations-table-quoted: %w", err) 165 } 166 } 167 if (len(migrationsTable) > 0) && (migrationsTableQuoted) && ((migrationsTable[0] != '"') || (migrationsTable[len(migrationsTable)-1] != '"')) { 168 return nil, fmt.Errorf("x-migrations-table must be quoted (for instance '\"migrate\".\"schema_migrations\"') when x-migrations-table-quoted is enabled, current value is: %s", migrationsTable) 169 } 170 171 statementTimeoutString := purl.Query().Get("x-statement-timeout") 172 statementTimeout := 0 173 if statementTimeoutString != "" { 174 statementTimeout, err = strconv.Atoi(statementTimeoutString) 175 if err != nil { 176 return nil, err 177 } 178 } 179 180 multiStatementMaxSize := DefaultMultiStatementMaxSize 181 if s := purl.Query().Get("x-multi-statement-max-size"); len(s) > 0 { 182 multiStatementMaxSize, err = strconv.Atoi(s) 183 if err != nil { 184 return nil, err 185 } 186 if multiStatementMaxSize <= 0 { 187 multiStatementMaxSize = DefaultMultiStatementMaxSize 188 } 189 } 190 191 multiStatementEnabled := false 192 if s := purl.Query().Get("x-multi-statement"); len(s) > 0 { 193 multiStatementEnabled, err = strconv.ParseBool(s) 194 if err != nil { 195 return nil, fmt.Errorf("Unable to parse option x-multi-statement: %w", err) 196 } 197 } 198 199 px, err := WithInstance(db, &Config{ 200 DatabaseName: purl.Path, 201 MigrationsTable: migrationsTable, 202 MigrationsTableQuoted: migrationsTableQuoted, 203 StatementTimeout: time.Duration(statementTimeout) * time.Millisecond, 204 MultiStatementEnabled: multiStatementEnabled, 205 MultiStatementMaxSize: multiStatementMaxSize, 206 }) 207 208 if err != nil { 209 return nil, err 210 } 211 212 return px, nil 213 } 214 215 func (p *Postgres) Close() error { 216 connErr := p.conn.Close() 217 dbErr := p.db.Close() 218 if connErr != nil || dbErr != nil { 219 return fmt.Errorf("conn: %v, db: %v", connErr, dbErr) 220 } 221 return nil 222 } 223 224 // https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS 225 func (p *Postgres) Lock() error { 226 return database.CasRestoreOnErr(&p.isLocked, false, true, database.ErrLocked, func() error { 227 aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName) 228 if err != nil { 229 return err 230 } 231 232 // This will wait indefinitely until the lock can be acquired. 233 query := `SELECT pg_advisory_lock($1)` 234 if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil { 235 return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)} 236 } 237 return nil 238 }) 239 } 240 241 func (p *Postgres) Unlock() error { 242 return database.CasRestoreOnErr(&p.isLocked, true, false, database.ErrNotLocked, func() error { 243 aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName) 244 if err != nil { 245 return err 246 } 247 248 query := `SELECT pg_advisory_unlock($1)` 249 if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil { 250 return &database.Error{OrigErr: err, Query: []byte(query)} 251 } 252 return nil 253 }) 254 } 255 256 func (p *Postgres) Run(migration io.Reader) error { 257 if p.config.MultiStatementEnabled { 258 var err error 259 if e := multistmt.Parse(migration, multiStmtDelimiter, p.config.MultiStatementMaxSize, func(m []byte) bool { 260 if err = p.runStatement(m); err != nil { 261 return false 262 } 263 return true 264 }); e != nil { 265 return e 266 } 267 return err 268 } 269 migr, err := io.ReadAll(migration) 270 if err != nil { 271 return err 272 } 273 return p.runStatement(migr) 274 } 275 276 func (p *Postgres) runStatement(statement []byte) error { 277 ctx := context.Background() 278 if p.config.StatementTimeout != 0 { 279 var cancel context.CancelFunc 280 ctx, cancel = context.WithTimeout(ctx, p.config.StatementTimeout) 281 defer cancel() 282 } 283 query := string(statement) 284 if strings.TrimSpace(query) == "" { 285 return nil 286 } 287 if _, err := p.conn.ExecContext(ctx, query); err != nil { 288 289 if pgErr, ok := err.(*pgconn.PgError); ok { 290 var line uint 291 var col uint 292 var lineColOK bool 293 line, col, lineColOK = computeLineFromPos(query, int(pgErr.Position)) 294 message := fmt.Sprintf("migration failed: %s", pgErr.Message) 295 if lineColOK { 296 message = fmt.Sprintf("%s (column %d)", message, col) 297 } 298 if pgErr.Detail != "" { 299 message = fmt.Sprintf("%s, %s", message, pgErr.Detail) 300 } 301 return database.Error{OrigErr: err, Err: message, Query: statement, Line: line} 302 } 303 return database.Error{OrigErr: err, Err: "migration failed", Query: statement} 304 } 305 return nil 306 } 307 308 func computeLineFromPos(s string, pos int) (line uint, col uint, ok bool) { 309 // replace crlf with lf 310 s = strings.Replace(s, "\r\n", "\n", -1) 311 // pg docs: pos uses index 1 for the first character, and positions are measured in characters not bytes 312 runes := []rune(s) 313 if pos > len(runes) { 314 return 0, 0, false 315 } 316 sel := runes[:pos] 317 line = uint(runesCount(sel, newLine) + 1) 318 col = uint(pos - 1 - runesLastIndex(sel, newLine)) 319 return line, col, true 320 } 321 322 const newLine = '\n' 323 324 func runesCount(input []rune, target rune) int { 325 var count int 326 for _, r := range input { 327 if r == target { 328 count++ 329 } 330 } 331 return count 332 } 333 334 func runesLastIndex(input []rune, target rune) int { 335 for i := len(input) - 1; i >= 0; i-- { 336 if input[i] == target { 337 return i 338 } 339 } 340 return -1 341 } 342 343 func (p *Postgres) SetVersion(version int, dirty bool) error { 344 tx, err := p.conn.BeginTx(context.Background(), &sql.TxOptions{}) 345 if err != nil { 346 return &database.Error{OrigErr: err, Err: "transaction start failed"} 347 } 348 349 query := `TRUNCATE ` + quoteIdentifier(p.config.migrationsSchemaName) + `.` + quoteIdentifier(p.config.migrationsTableName) 350 if _, err := tx.Exec(query); err != nil { 351 if errRollback := tx.Rollback(); errRollback != nil { 352 err = multierror.Append(err, errRollback) 353 } 354 return &database.Error{OrigErr: err, Query: []byte(query)} 355 } 356 357 // Also re-write the schema version for nil dirty versions to prevent 358 // empty schema version for failed down migration on the first migration 359 // See: https://github.com/golang-migrate/migrate/issues/330 360 if version >= 0 || (version == database.NilVersion && dirty) { 361 query = `INSERT INTO ` + quoteIdentifier(p.config.migrationsSchemaName) + `.` + quoteIdentifier(p.config.migrationsTableName) + ` (version, dirty) VALUES ($1, $2)` 362 if _, err := tx.Exec(query, version, dirty); err != nil { 363 if errRollback := tx.Rollback(); errRollback != nil { 364 err = multierror.Append(err, errRollback) 365 } 366 return &database.Error{OrigErr: err, Query: []byte(query)} 367 } 368 } 369 370 if err := tx.Commit(); err != nil { 371 return &database.Error{OrigErr: err, Err: "transaction commit failed"} 372 } 373 374 return nil 375 } 376 377 func (p *Postgres) Version() (version int, dirty bool, err error) { 378 query := `SELECT version, dirty FROM ` + quoteIdentifier(p.config.migrationsSchemaName) + `.` + quoteIdentifier(p.config.migrationsTableName) + ` LIMIT 1` 379 err = p.conn.QueryRowContext(context.Background(), query).Scan(&version, &dirty) 380 switch { 381 case err == sql.ErrNoRows: 382 return database.NilVersion, false, nil 383 384 case err != nil: 385 if e, ok := err.(*pgconn.PgError); ok { 386 if e.SQLState() == pgerrcode.UndefinedTable { 387 return database.NilVersion, false, nil 388 } 389 } 390 return 0, false, &database.Error{OrigErr: err, Query: []byte(query)} 391 392 default: 393 return version, dirty, nil 394 } 395 } 396 397 func (p *Postgres) Drop() (err error) { 398 // select all tables in current schema 399 query := `SELECT table_name FROM information_schema.tables WHERE table_schema=(SELECT current_schema()) AND table_type='BASE TABLE'` 400 tables, err := p.conn.QueryContext(context.Background(), query) 401 if err != nil { 402 return &database.Error{OrigErr: err, Query: []byte(query)} 403 } 404 defer func() { 405 if errClose := tables.Close(); errClose != nil { 406 err = multierror.Append(err, errClose) 407 } 408 }() 409 410 // delete one table after another 411 tableNames := make([]string, 0) 412 for tables.Next() { 413 var tableName string 414 if err := tables.Scan(&tableName); err != nil { 415 return err 416 } 417 if len(tableName) > 0 { 418 tableNames = append(tableNames, tableName) 419 } 420 } 421 if err := tables.Err(); err != nil { 422 return &database.Error{OrigErr: err, Query: []byte(query)} 423 } 424 425 if len(tableNames) > 0 { 426 // delete one by one ... 427 for _, t := range tableNames { 428 query = `DROP TABLE IF EXISTS ` + quoteIdentifier(t) + ` CASCADE` 429 if _, err := p.conn.ExecContext(context.Background(), query); err != nil { 430 return &database.Error{OrigErr: err, Query: []byte(query)} 431 } 432 } 433 } 434 435 return nil 436 } 437 438 // ensureVersionTable checks if versions table exists and, if not, creates it. 439 // Note that this function locks the database, which deviates from the usual 440 // convention of "caller locks" in the Postgres type. 441 func (p *Postgres) ensureVersionTable() (err error) { 442 if err = p.Lock(); err != nil { 443 return err 444 } 445 446 defer func() { 447 if e := p.Unlock(); e != nil { 448 if err == nil { 449 err = e 450 } else { 451 err = multierror.Append(err, e) 452 } 453 } 454 }() 455 456 // This block checks whether the `MigrationsTable` already exists. This is useful because it allows read only postgres 457 // users to also check the current version of the schema. Previously, even if `MigrationsTable` existed, the 458 // `CREATE TABLE IF NOT EXISTS...` query would fail because the user does not have the CREATE permission. 459 // Taken from https://github.com/mattes/migrate/blob/master/database/postgres/postgres.go#L258 460 query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_schema = $1 AND table_name = $2 LIMIT 1` 461 row := p.conn.QueryRowContext(context.Background(), query, p.config.migrationsSchemaName, p.config.migrationsTableName) 462 463 var count int 464 err = row.Scan(&count) 465 if err != nil { 466 return &database.Error{OrigErr: err, Query: []byte(query)} 467 } 468 469 if count == 1 { 470 return nil 471 } 472 473 query = `CREATE TABLE IF NOT EXISTS ` + quoteIdentifier(p.config.migrationsSchemaName) + `.` + quoteIdentifier(p.config.migrationsTableName) + ` (version bigint not null primary key, dirty boolean not null)` 474 if _, err = p.conn.ExecContext(context.Background(), query); err != nil { 475 return &database.Error{OrigErr: err, Query: []byte(query)} 476 } 477 478 return nil 479 } 480 481 // Copied from lib/pq implementation: https://github.com/lib/pq/blob/v1.9.0/conn.go#L1611 482 func quoteIdentifier(name string) string { 483 end := strings.IndexRune(name, 0) 484 if end > -1 { 485 name = name[:end] 486 } 487 return `"` + strings.Replace(name, `"`, `""`, -1) + `"` 488 }