storj/private/dbutil/cockroachutil/db.go

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.

package cockroachutil

import (
	"database/sql"
	"net/url"
	"strings"

	"github.com/lib/pq"
	"github.com/zeebo/errs"
	"gopkg.in/spacemonkeygo/monkit.v2"

	"storj.io/storj/private/dbutil"
)

var mon = monkit.Package()

// 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, schemaName string) (db *dbutil.TempDatabase, err error) {
	if !strings.HasPrefix(connStr, "cockroach://") {
		return nil, errs.New("expected a cockroachDB URI, but got %q", connStr)
	}
	postgresConnStr := "postgres://" + connStr[12:]
	masterDB, err := sql.Open("postgres", postgresConnStr)
	if err != nil {
		return nil, errs.Wrap(err)
	}
	defer func() {
		err = errs.Combine(err, masterDB.Close())
	}()
	err = masterDB.Ping()
	if err != nil {
		return nil, errs.New("Could not open masterDB at conn %q: %v", connStr, err)
	}

	_, err = masterDB.Exec("CREATE DATABASE " + pq.QuoteIdentifier(schemaName))
	if err != nil {
		return nil, errs.Wrap(err)
	}

	cleanup := func(cleanupDB *sql.DB) error {
		_, err := cleanupDB.Exec("DROP DATABASE " + pq.QuoteIdentifier(schemaName))
		return errs.Wrap(err)
	}

	modifiedConnStr, err := changeDBTargetInConnStr(connStr, schemaName)
	if err != nil {
		return nil, errs.Combine(err, cleanup(masterDB))
	}

	modifiedPostgresConnStr := "postgres://" + modifiedConnStr[12:]
	sqlDB, err := sql.Open("postgres", modifiedPostgresConnStr)
	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:         "postgres",
		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
}