1279eeae39
Replace all the remaining uses of sql.DB with tagsql.DB to fix issues with context cancellation. Introduce tagsql.Open which helps to get rid of all tagsql.Wrap-s. Use tagsql in cockroachkv and postgreskv. Change-Id: I8946d203341cb85a25976896fc7881e1f704e779
94 lines
2.5 KiB
Go
94 lines
2.5 KiB
Go
// 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
|
|
}
|