storj/private/dbutil/cockroachutil/db.go

94 lines
2.5 KiB
Go
Raw Normal View History

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.
package cockroachutil
import (
"context"
"encoding/hex"
"math/rand"
"net/url"
"strings"
"github.com/lib/pq"
"github.com/zeebo/errs"
"gopkg.in/spacemonkeygo/monkit.v2"
"storj.io/storj/private/dbutil"
"storj.io/storj/private/tagsql"
)
var mon = monkit.Package()
// CreateRandomTestingSchemaName creates a random schema name string.
func CreateRandomTestingSchemaName(n int) string {
data := make([]byte, n)
// math/rand.Read() always returns a nil error so there's no need to handle the error.
_, _ = rand.Read(data)
return hex.EncodeToString(data)
}
// OpenUnique opens a temporary unique CockroachDB database that will be cleaned up when closed.
// It is expected that this should normally be used by way of
// "storj.io/storj/private/dbutil/tempdb".OpenUnique() instead of calling it directly.
func OpenUnique(ctx context.Context, connStr string, schemaPrefix string) (db *dbutil.TempDatabase, err error) {
if !strings.HasPrefix(connStr, "cockroach://") {
return nil, errs.New("expected a cockroachDB URI, but got %q", connStr)
}
schemaName := schemaPrefix + "-" + CreateRandomTestingSchemaName(8)
masterDB, err := tagsql.Open("cockroach", connStr)
if err != nil {
return nil, errs.Wrap(err)
}
defer func() {
err = errs.Combine(err, masterDB.Close())
}()
err = masterDB.PingContext(ctx)
if err != nil {
return nil, errs.New("Could not open masterDB at conn %q: %w", connStr, err)
}
_, err = masterDB.Exec(ctx, "CREATE DATABASE "+pq.QuoteIdentifier(schemaName))
if err != nil {
return nil, errs.Wrap(err)
}
cleanup := func(cleanupDB tagsql.DB) error {
_, err := cleanupDB.Exec(context.TODO(), "DROP DATABASE "+pq.QuoteIdentifier(schemaName))
return errs.Wrap(err)
}
modifiedConnStr, err := changeDBTargetInConnStr(connStr, schemaName)
if err != nil {
return nil, errs.Combine(err, cleanup(masterDB))
}
sqlDB, err := tagsql.Open("cockroach", modifiedConnStr)
if err != nil {
return nil, errs.Combine(errs.Wrap(err), cleanup(masterDB))
}
dbutil.Configure(sqlDB, mon)
return &dbutil.TempDatabase{
DB: sqlDB,
ConnStr: modifiedConnStr,
Schema: schemaName,
Driver: "cockroach",
Implementation: dbutil.Cockroach,
Cleanup: cleanup,
}, nil
}
func changeDBTargetInConnStr(connStr string, newDBName string) (string, error) {
connURL, err := url.Parse(connStr)
if err != nil {
return "", errs.Wrap(err)
}
connURL.Path = newDBName
return connURL.String(), nil
}