storj/private/dbutil/cockroachutil/db.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
}