From b9740f0c0a2d7daec5a12e94bf96cec6d5ac04fa Mon Sep 17 00:00:00 2001
From: Egon Elbre <egonelbre@gmail.com>
Date: Mon, 13 Jan 2020 15:57:47 +0200
Subject: [PATCH] storage/cockroachkv: add ctx argument

Change-Id: Ib6c29f44722b0354afcd499a0e567f04aef7eb28
---
 private/dbutil/cockroachutil/db.go    |  3 ++-
 private/dbutil/tempdb/tempdb.go       |  2 +-
 storage/cockroachkv/client.go         |  3 ++-
 storage/cockroachkv/client_test.go    | 13 +++++++++----
 storage/cockroachkv/schema/migrate.go |  3 ++-
 5 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/private/dbutil/cockroachutil/db.go b/private/dbutil/cockroachutil/db.go
index 6955e4f7f..4da847f2f 100644
--- a/private/dbutil/cockroachutil/db.go
+++ b/private/dbutil/cockroachutil/db.go
@@ -4,6 +4,7 @@
 package cockroachutil
 
 import (
+	"context"
 	"database/sql"
 	"encoding/hex"
 	"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.
 // 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(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://") {
 		return nil, errs.New("expected a cockroachDB URI, but got %q", connStr)
 	}
diff --git a/private/dbutil/tempdb/tempdb.go b/private/dbutil/tempdb/tempdb.go
index 1264cb9ea..c7eb2c782 100644
--- a/private/dbutil/tempdb/tempdb.go
+++ b/private/dbutil/tempdb/tempdb.go
@@ -21,7 +21,7 @@ func OpenUnique(ctx context.Context, connURL string, namePrefix string) (*dbutil
 		return pgutil.OpenUnique(ctx, connURL, namePrefix)
 	}
 	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)
 }
diff --git a/storage/cockroachkv/client.go b/storage/cockroachkv/client.go
index 14b6f4f05..d69d32b48 100644
--- a/storage/cockroachkv/client.go
+++ b/storage/cockroachkv/client.go
@@ -44,7 +44,8 @@ func New(dbURL string) (*Client, error) {
 
 	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 {
 		return nil, err
 	}
diff --git a/storage/cockroachkv/client_test.go b/storage/cockroachkv/client_test.go
index 199605f24..b2631f4f9 100644
--- a/storage/cockroachkv/client_test.go
+++ b/storage/cockroachkv/client_test.go
@@ -3,27 +3,29 @@
 package cockroachkv
 
 import (
+	"context"
 	"testing"
 
 	_ "github.com/lib/pq"
 
+	"storj.io/common/testcontext"
 	"storj.io/storj/private/dbutil/cockroachutil"
 	"storj.io/storj/private/dbutil/pgutil/pgtest"
 	"storj.io/storj/storage/cockroachkv/schema"
 	"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 == "" {
 		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 {
 		t.Fatalf("init: %+v", err)
 	}
 
-	err = schema.PrepareDB(tdb.DB)
+	err = schema.PrepareDB(ctx, tdb.DB)
 	if err != nil {
 		t.Fatalf("init: %+v", err)
 	}
@@ -36,7 +38,10 @@ func newTestCockroachDB(t testing.TB) (store *Client, cleanup func()) {
 }
 
 func TestSuite(t *testing.T) {
-	store, cleanup := newTestCockroachDB(t)
+	ctx := testcontext.New(t)
+	defer ctx.Cleanup()
+
+	store, cleanup := newTestCockroachDB(ctx, t)
 	defer cleanup()
 
 	testsuite.RunTests(t, store)
diff --git a/storage/cockroachkv/schema/migrate.go b/storage/cockroachkv/schema/migrate.go
index c30084e99..46361b364 100644
--- a/storage/cockroachkv/schema/migrate.go
+++ b/storage/cockroachkv/schema/migrate.go
@@ -4,6 +4,7 @@
 package schema
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 
@@ -12,7 +13,7 @@ import (
 )
 
 // 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
 	if err := db.QueryRow(`SELECT current_database();`).Scan(&dbName); err != nil {
 		return errs.Wrap(err)