github.com/status-im/status-go@v1.1.0/postgres/helpers.go (about)

     1  package postgres
     2  
     3  import (
     4  	"database/sql"
     5  	"fmt"
     6  	"os"
     7  
     8  	// Import postgres driver
     9  	_ "github.com/lib/pq"
    10  )
    11  
    12  var (
    13  	DefaultTestDBHost = GetEnvDefault("TEST_POSTGRES_HOST", "localhost")
    14  	DefaultTestDBPort = GetEnvDefault("TEST_POSTGRES_PORT", "5432")
    15  	DefaultTestURI    = fmt.Sprintf("postgres://postgres@%s:%s/postgres?sslmode=disable", DefaultTestDBHost, DefaultTestDBPort)
    16  	DropTableURI      = fmt.Sprintf("postgres://postgres@%s:%s/template1?sslmode=disable", DefaultTestDBHost, DefaultTestDBPort)
    17  )
    18  
    19  func GetEnvDefault(key, fallback string) string {
    20  	value := os.Getenv(key)
    21  	if len(value) == 0 {
    22  		return fallback
    23  	}
    24  	return value
    25  }
    26  
    27  func ResetDefaultTestPostgresDB() error {
    28  	db, err := sql.Open("postgres", DropTableURI)
    29  	if err != nil {
    30  		return err
    31  	}
    32  	defer func() {
    33  		_ = db.Close()
    34  	}()
    35  
    36  	// Drop current and prevent any future connections. Used in tests. Details here:
    37  	// https://stackoverflow.com/questions/17449420/postgresql-unable-to-drop-database-because-of-some-auto-connections-to-db
    38  	_, err = db.Exec("REVOKE CONNECT ON DATABASE postgres FROM public;")
    39  	if err != nil {
    40  		return err
    41  	}
    42  
    43  	_, err = db.Exec("SELECT pid, pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = 'postgres' AND pid <> pg_backend_pid();")
    44  	if err != nil {
    45  		return err
    46  	}
    47  
    48  	_, err = db.Exec("DROP DATABASE IF EXISTS postgres;")
    49  	if err != nil {
    50  		return err
    51  	}
    52  
    53  	_, err = db.Exec("CREATE DATABASE postgres;")
    54  	if err != nil {
    55  		return err
    56  	}
    57  
    58  	_, err = db.Exec("GRANT CONNECT ON DATABASE postgres TO public;")
    59  	if err != nil {
    60  		return err
    61  	}
    62  
    63  	return nil
    64  }