database/pg/pgtest/pgtest.go

   1 package pgtest
   2
   3 import (
   4         "context"
   5         "database/sql"
   6         "io/ioutil"
   7         "log"
   8         "math/rand"
   9         "net/url"
  10         "os"
  11         "runtime"
  12         "testing"
  13         "time"
  14
  15         "github.com/lib/pq"
  16
  17         "chain/database/pg"
  18         "chain/testutil"
  19 )
  20
  21 var (
  22         random = rand.New(rand.NewSource(time.Now().UnixNano()))
  23
  24         // dbpool contains initialized, pristine databases,
  25         // as returned from open. It is the client's job to
  26         // make sure a database is in this state
  27         // (for example, by rolling back a transaction)
  28         // before returning it to the pool.
  29         dbpool = make(chan *sql.DB, 4)
  30 )
  31
  32 // DefaultURL is used by NewTX and NewDB if DBURL is the empty string.
  33 const DefaultURL = "postgres:///postgres?sslmode=disable"
  34
  35 var (
  36         // DBURL should be a URL of the form "postgres://...".
  37         // If it is the empty string, DefaultURL will be used.
  38         // The functions NewTx and NewDB use it to create and connect
  39         // to new databases by replacing the database name component
  40         // with a randomized name.
  41         DBURL = os.Getenv("DB_URL_TEST")
  42
  43         // SchemaPath is a file containing a schema to initialize
  44         // a database in NewTx.
  45         SchemaPath = os.Getenv("CHAIN") + "/core/schema.sql"
  46 )
  47
  48 const (
  49         gcDur      = 3 * time.Minute
  50         timeFormat = "20060102150405"
  51 )
  52
  53 // NewDB creates a database initialized
  54 // with the schema in schemaPath.
  55 // It returns the resulting *sql.DB with its URL.
  56 //
  57 // It also registers a finalizer for the DB, so callers
  58 // can discard it without closing it explicitly, and the
  59 // test program is nevertheless unlikely to run out of
  60 // connection slots in the server.
  61 //
  62 // Prefer NewTx whenever the caller can do its
  63 // work in exactly one transaction.
  64 func NewDB(f Fataler, schemaPath string) (url string, db *sql.DB) {
  65         ctx := context.Background()
  66         if os.Getenv("CHAIN") == "" {
  67                 log.Println("warning: $CHAIN not set; probably can't find schema")
  68         }
  69         url, db, err := open(ctx, DBURL, schemaPath)
  70         if err != nil {
  71                 f.Fatal(err)
  72         }
  73         runtime.SetFinalizer(db, (*sql.DB).Close)
  74         return url, db
  75 }
  76
  77 // NewTx returns a new transaction on a database
  78 // initialized with the schema in SchemaPath.
  79 //
  80 // It also registers a finalizer for the Tx, so callers
  81 // can discard it without rolling back explicitly, and the
  82 // test program is nevertheless unlikely to run out of
  83 // connection slots in the server.
  84 // The caller should not commit the returned Tx; doing so
  85 // will prevent the underlying database from being reused
  86 // and so cause future calls to NewTx to be slower.
  87 func NewTx(f Fataler) *sql.Tx {
  88         runtime.GC() // give the finalizers a better chance to run
  89         ctx := context.Background()
  90         if os.Getenv("CHAIN") == "" {
  91                 log.Println("warning: $CHAIN not set; probably can't find schema")
  92         }
  93         db, err := getdb(ctx, DBURL, SchemaPath)
  94         if err != nil {
  95                 f.Fatal(err)
  96         }
  97         tx, err := db.BeginTx(ctx, nil)
  98         if err != nil {
  99                 db.Close()
 100                 f.Fatal(err)
 101         }
 102         // NOTE(kr): we do not set a finalizer on the DB.
 103         // It is closed explicitly, if necessary, by finalizeTx.
 104         runtime.SetFinalizer(tx, finaldb{db}.finalizeTx)
 105         return tx
 106 }
 107
 108 // CloneDB creates a new database, using the database at the provided
 109 // URL as a template. It returns the URL of the database clone.
 110 func CloneDB(ctx context.Context, baseURL string) (newURL string, err error) {
 111         u, err := url.Parse(baseURL)
 112         if err != nil {
 113                 return "", err
 114         }
 115
 116         ctldb, err := sql.Open("postgres", baseURL)
 117         if err != nil {
 118                 return "", err
 119         }
 120         defer ctldb.Close()
 121
 122         dbname := pickName("db")
 123         _, err = ctldb.Exec("CREATE DATABASE " + pq.QuoteIdentifier(dbname) + " WITH TEMPLATE " + pq.QuoteIdentifier(u.Path[1:]))
 124         if err != nil {
 125                 return "", err
 126         }
 127         u.Path = "/" + dbname
 128         return u.String(), nil
 129 }
 130
 131 // open derives a new randomized test database name from baseURL,
 132 // initializes it with schemaFile, and opens it.
 133 func open(ctx context.Context, baseURL, schemaFile string) (newurl string, db *sql.DB, err error) {
 134         if baseURL == "" {
 135                 baseURL = DefaultURL
 136         }
 137
 138         u, err := url.Parse(baseURL)
 139         if err != nil {
 140                 return "", nil, err
 141         }
 142
 143         ctldb, err := sql.Open("postgres", baseURL)
 144         if err != nil {
 145                 return "", nil, err
 146         }
 147         defer ctldb.Close()
 148
 149         err = gcdbs(ctldb)
 150         if err != nil {
 151                 log.Println(err)
 152         }
 153
 154         dbname := pickName("db")
 155         u.Path = "/" + dbname
 156         _, err = ctldb.Exec("CREATE DATABASE " + pq.QuoteIdentifier(dbname))
 157         if err != nil {
 158                 return "", nil, err
 159         }
 160
 161         schema, err := ioutil.ReadFile(schemaFile)
 162         if err != nil {
 163                 return "", nil, err
 164         }
 165         db, err = sql.Open("postgres", u.String())
 166         if err != nil {
 167                 return "", nil, err
 168         }
 169         _, err = db.ExecContext(ctx, string(schema))
 170         if err != nil {
 171                 db.Close()
 172                 return "", nil, err
 173         }
 174         return u.String(), db, nil
 175 }
 176
 177 type finaldb struct{ db *sql.DB }
 178
 179 func (f finaldb) finalizeTx(tx *sql.Tx) {
 180         go func() { // don't block the finalizer goroutine for too long
 181                 err := tx.Rollback()
 182                 if err != nil {
 183                         // If the tx has been committed (or if anything
 184                         // else goes wrong), we can't reuse db.
 185                         f.db.Close()
 186                         return
 187                 }
 188                 select {
 189                 case dbpool <- f.db:
 190                 default:
 191                         f.db.Close() // pool is full
 192                 }
 193         }()
 194 }
 195
 196 func getdb(ctx context.Context, url, path string) (*sql.DB, error) {
 197         select {
 198         case db := <-dbpool:
 199                 return db, nil
 200         default:
 201                 _, db, err := open(ctx, url, path)
 202                 return db, err
 203         }
 204 }
 205
 206 func gcdbs(db *sql.DB) error {
 207         gcTime := time.Now().Add(-gcDur)
 208         const q = `
 209                 SELECT datname FROM pg_database
 210                 WHERE datname LIKE 'pgtest_%' AND datname < $1
 211         `
 212         rows, err := db.Query(q, formatPrefix("db", gcTime))
 213         if err != nil {
 214                 return err
 215         }
 216         var names []string
 217         for rows.Next() {
 218                 var name string
 219                 err = rows.Scan(&name)
 220                 if err != nil {
 221                         return err
 222                 }
 223                 names = append(names, name)
 224         }
 225         if rows.Err() != nil {
 226                 return rows.Err()
 227         }
 228         for i, name := range names {
 229                 if i > 5 {
 230                         break // drop up to five databases per test
 231                 }
 232                 go db.Exec("DROP DATABASE " + pq.QuoteIdentifier(name))
 233         }
 234         return nil
 235 }
 236
 237 func pickName(prefix string) (s string) {
 238         const chars = "abcdefghijklmnopqrstuvwxyz"
 239         for i := 0; i < 10; i++ {
 240                 s += string(chars[random.Intn(len(chars))])
 241         }
 242         return formatPrefix(prefix, time.Now()) + s
 243 }
 244
 245 func formatPrefix(prefix string, t time.Time) string {
 246         return "pgtest_" + prefix + "_" + t.UTC().Format(timeFormat) + "Z_"
 247 }
 248
 249 // Exec executes q in the database or transaction in ctx.
 250 // If there is an error, it fails t.
 251 func Exec(ctx context.Context, db pg.DB, t testing.TB, q string, args ...interface{}) {
 252         _, err := db.ExecContext(ctx, q, args...)
 253         if err != nil {
 254                 testutil.FatalErr(t, err)
 255         }
 256 }
 257
 258 // Fataler lets NewTx and NewDB signal immediate failure.
 259 // It is satisfied by *testing.T, *testing.B, and *log.Logger.
 260 type Fataler interface {
 261         Fatal(...interface{})
 262 }