golang.org/x/build@v0.0.0-20240506185731-218518f32b70/internal/relui/store.go (about)

     1  // Copyright 2020 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package relui
     6  
     7  import (
     8  	"context"
     9  	"database/sql"
    10  	"errors"
    11  	"fmt"
    12  
    13  	"github.com/golang-migrate/migrate/v4"
    14  	dbpgx "github.com/golang-migrate/migrate/v4/database/pgx"
    15  	"github.com/golang-migrate/migrate/v4/source/iofs"
    16  	"github.com/jackc/pgx/v4"
    17  )
    18  
    19  var errDBNotExist = errors.New("database does not exist")
    20  
    21  // InitDB creates and applies all migrations to the database specified
    22  // in conn.
    23  //
    24  // If the database does not exist, one will be created using the
    25  // credentials provided.
    26  //
    27  // Any key/value or URI string compatible with libpq is valid.
    28  func InitDB(ctx context.Context, conn string) error {
    29  	cfg, err := pgx.ParseConfig(conn)
    30  	if err != nil {
    31  		return fmt.Errorf("pgx.ParseConfig() = %w", err)
    32  	}
    33  	if err := CreateDBIfNotExists(ctx, cfg); err != nil {
    34  		return err
    35  	}
    36  	if err := MigrateDB(conn, false); err != nil {
    37  		return err
    38  	}
    39  	return nil
    40  }
    41  
    42  // MigrateDB applies all migrations to the database specified in conn.
    43  //
    44  // Any key/value or URI string compatible with libpq is a valid conn.
    45  // If downUp is true, all migrations will be run, then the down and up
    46  // migrations of the final migration are run.
    47  func MigrateDB(conn string, downUp bool) error {
    48  	cfg, err := pgx.ParseConfig(conn)
    49  	if err != nil {
    50  		return fmt.Errorf("pgx.ParseConfig() = %w", err)
    51  	}
    52  	db, err := sql.Open("pgx", conn)
    53  	if err != nil {
    54  		return fmt.Errorf("sql.Open(%q, _) = %v, %w", "pgx", db, err)
    55  	}
    56  	mcfg := &dbpgx.Config{
    57  		MigrationsTable: "migrations",
    58  		DatabaseName:    cfg.Database,
    59  	}
    60  	mdb, err := dbpgx.WithInstance(db, mcfg)
    61  	if err != nil {
    62  		return fmt.Errorf("dbpgx.WithInstance(_, %v) = %v, %w", mcfg, mdb, err)
    63  	}
    64  	mfs, err := iofs.New(migrations, "migrations")
    65  	if err != nil {
    66  		return fmt.Errorf("iofs.New(%v, %q) = %v, %w", migrations, "migrations", mfs, err)
    67  	}
    68  	m, err := migrate.NewWithInstance("iofs", mfs, "pgx", mdb)
    69  	if err != nil {
    70  		return fmt.Errorf("migrate.NewWithInstance(%q, %v, %q, %v) = %v, %w", "iofs", migrations, "pgx", mdb, m, err)
    71  	}
    72  	if err := m.Up(); err != nil && !errors.Is(err, migrate.ErrNoChange) {
    73  		return fmt.Errorf("m.Up() = %w", err)
    74  	}
    75  	if downUp {
    76  		if err := m.Steps(-1); err != nil {
    77  			return fmt.Errorf("m.Steps(%d) = %w", -1, err)
    78  		}
    79  		if err := m.Up(); err != nil {
    80  			return fmt.Errorf("m.Up() = %w", err)
    81  		}
    82  	}
    83  	db.Close()
    84  	return nil
    85  }
    86  
    87  // ConnectMaintenanceDB connects to the maintenance database using the
    88  // credentials from cfg. If maintDB is an empty string, the database
    89  // with the name cfg.User will be used.
    90  func ConnectMaintenanceDB(ctx context.Context, cfg *pgx.ConnConfig, maintDB string) (*pgx.Conn, error) {
    91  	cfg = cfg.Copy()
    92  	if maintDB == "" {
    93  		maintDB = "postgres"
    94  	}
    95  	cfg.Database = maintDB
    96  	return pgx.ConnectConfig(ctx, cfg)
    97  }
    98  
    99  // CreateDBIfNotExists checks whether the given dbName is an existing
   100  // database, and creates one if not.
   101  func CreateDBIfNotExists(ctx context.Context, cfg *pgx.ConnConfig) error {
   102  	exists, err := checkIfDBExists(ctx, cfg)
   103  	if err != nil || exists {
   104  		return err
   105  	}
   106  	conn, err := ConnectMaintenanceDB(ctx, cfg, "")
   107  	if err != nil {
   108  		return fmt.Errorf("ConnectMaintenanceDB = %w", err)
   109  	}
   110  	createSQL := fmt.Sprintf("CREATE DATABASE %s", pgx.Identifier{cfg.Database}.Sanitize())
   111  	if _, err := conn.Exec(ctx, createSQL); err != nil {
   112  		return fmt.Errorf("conn.Exec(%q) = %w", createSQL, err)
   113  	}
   114  	return nil
   115  }
   116  
   117  // DropDB drops the database specified in cfg. An error returned if
   118  // the database does not exist.
   119  func DropDB(ctx context.Context, cfg *pgx.ConnConfig) error {
   120  	exists, err := checkIfDBExists(ctx, cfg)
   121  	if err != nil {
   122  		return fmt.Errorf("p.checkIfDBExists() = %w", err)
   123  	}
   124  	if !exists {
   125  		return errDBNotExist
   126  	}
   127  	conn, err := ConnectMaintenanceDB(ctx, cfg, "")
   128  	if err != nil {
   129  		return fmt.Errorf("ConnectMaintenanceDB = %w", err)
   130  	}
   131  	dropSQL := fmt.Sprintf("DROP DATABASE %s", pgx.Identifier{cfg.Database}.Sanitize())
   132  	if _, err := conn.Exec(ctx, dropSQL); err != nil {
   133  		return fmt.Errorf("conn.Exec(%q) = %w", dropSQL, err)
   134  	}
   135  	return nil
   136  }
   137  
   138  func checkIfDBExists(ctx context.Context, cfg *pgx.ConnConfig) (bool, error) {
   139  	conn, err := ConnectMaintenanceDB(ctx, cfg, "")
   140  	if err != nil {
   141  		return false, fmt.Errorf("ConnectMaintenanceDB = %w", err)
   142  	}
   143  	row := conn.QueryRow(ctx, "SELECT 1 from pg_database WHERE datname=$1 LIMIT 1", cfg.Database)
   144  	var exists int
   145  	if err := row.Scan(&exists); err != nil && err != pgx.ErrNoRows {
   146  		return false, fmt.Errorf("row.Scan() = %w", err)
   147  	}
   148  	return exists == 1, nil
   149  }