storage/cockroachkv: add ctx argument

Change-Id: Ib6c29f44722b0354afcd499a0e567f04aef7eb28
This commit is contained in:
Egon Elbre 2020-01-13 15:57:47 +02:00
parent ff267168c5
commit b9740f0c0a
5 changed files with 16 additions and 8 deletions

View File

@ -4,6 +4,7 @@
package cockroachutil package cockroachutil
import ( import (
"context"
"database/sql" "database/sql"
"encoding/hex" "encoding/hex"
"math/rand" "math/rand"
@ -31,7 +32,7 @@ func CreateRandomTestingSchemaName(n int) string {
// OpenUnique opens a temporary unique CockroachDB database that will be cleaned up when closed. // 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 // It is expected that this should normally be used by way of
// "storj.io/storj/private/dbutil/tempdb".OpenUnique() instead of calling it directly. // "storj.io/storj/private/dbutil/tempdb".OpenUnique() instead of calling it directly.
func OpenUnique(connStr string, schemaPrefix string) (db *dbutil.TempDatabase, err error) { func OpenUnique(ctx context.Context, connStr string, schemaPrefix string) (db *dbutil.TempDatabase, err error) {
if !strings.HasPrefix(connStr, "cockroach://") { if !strings.HasPrefix(connStr, "cockroach://") {
return nil, errs.New("expected a cockroachDB URI, but got %q", connStr) return nil, errs.New("expected a cockroachDB URI, but got %q", connStr)
} }

View File

@ -21,7 +21,7 @@ func OpenUnique(ctx context.Context, connURL string, namePrefix string) (*dbutil
return pgutil.OpenUnique(ctx, connURL, namePrefix) return pgutil.OpenUnique(ctx, connURL, namePrefix)
} }
if strings.HasPrefix(connURL, "cockroach://") { if strings.HasPrefix(connURL, "cockroach://") {
return cockroachutil.OpenUnique(connURL, namePrefix) return cockroachutil.OpenUnique(ctx, connURL, namePrefix)
} }
return nil, errs.New("OpenUnique does not yet support the db type for %q", connURL) return nil, errs.New("OpenUnique does not yet support the db type for %q", connURL)
} }

View File

@ -44,7 +44,8 @@ func New(dbURL string) (*Client, error) {
dbutil.Configure(db, mon) dbutil.Configure(db, mon)
err = schema.PrepareDB(db) // TODO: new shouldn't be taking ctx as argument
err = schema.PrepareDB(context.TODO(), db)
if err != nil { if err != nil {
return nil, err return nil, err
} }

View File

@ -3,27 +3,29 @@
package cockroachkv package cockroachkv
import ( import (
"context"
"testing" "testing"
_ "github.com/lib/pq" _ "github.com/lib/pq"
"storj.io/common/testcontext"
"storj.io/storj/private/dbutil/cockroachutil" "storj.io/storj/private/dbutil/cockroachutil"
"storj.io/storj/private/dbutil/pgutil/pgtest" "storj.io/storj/private/dbutil/pgutil/pgtest"
"storj.io/storj/storage/cockroachkv/schema" "storj.io/storj/storage/cockroachkv/schema"
"storj.io/storj/storage/testsuite" "storj.io/storj/storage/testsuite"
) )
func newTestCockroachDB(t testing.TB) (store *Client, cleanup func()) { func newTestCockroachDB(ctx context.Context, t testing.TB) (store *Client, cleanup func()) {
if *pgtest.CrdbConnStr == "" { if *pgtest.CrdbConnStr == "" {
t.Skipf("cockroach flag missing, example:\n-cockroach-test-db=%s", pgtest.DefaultCrdbConnStr) t.Skipf("cockroach flag missing, example:\n-cockroach-test-db=%s", pgtest.DefaultCrdbConnStr)
} }
tdb, err := cockroachutil.OpenUnique(*pgtest.CrdbConnStr, "test-schema") tdb, err := cockroachutil.OpenUnique(ctx, *pgtest.CrdbConnStr, "test-schema")
if err != nil { if err != nil {
t.Fatalf("init: %+v", err) t.Fatalf("init: %+v", err)
} }
err = schema.PrepareDB(tdb.DB) err = schema.PrepareDB(ctx, tdb.DB)
if err != nil { if err != nil {
t.Fatalf("init: %+v", err) t.Fatalf("init: %+v", err)
} }
@ -36,7 +38,10 @@ func newTestCockroachDB(t testing.TB) (store *Client, cleanup func()) {
} }
func TestSuite(t *testing.T) { func TestSuite(t *testing.T) {
store, cleanup := newTestCockroachDB(t) ctx := testcontext.New(t)
defer ctx.Cleanup()
store, cleanup := newTestCockroachDB(ctx, t)
defer cleanup() defer cleanup()
testsuite.RunTests(t, store) testsuite.RunTests(t, store)

View File

@ -4,6 +4,7 @@
package schema package schema
import ( import (
"context"
"database/sql" "database/sql"
"fmt" "fmt"
@ -12,7 +13,7 @@ import (
) )
// PrepareDB creates the pathdata tables if they don't already exist. // PrepareDB creates the pathdata tables if they don't already exist.
func PrepareDB(db *sql.DB) (err error) { func PrepareDB(ctx context.Context, db *sql.DB) (err error) {
var dbName string var dbName string
if err := db.QueryRow(`SELECT current_database();`).Scan(&dbName); err != nil { if err := db.QueryRow(`SELECT current_database();`).Scan(&dbName); err != nil {
return errs.Wrap(err) return errs.Wrap(err)