storj/private/dbutil/cockroachutil/driver_test.go
paul cannon 5a1838bc28 private/dbutil: retry single statements on cockroachdb
This ought to make it so that all single statements (Exec- or Query-) on
a CockroachDB backend will get retried as necessary. As there is no need
for savepoints to be allocated or released in this case, there is no
round-trip overhead except when statements actually do need to be
retried.

Change-Id: Ibd7f1725ff727477c456cb309120d080f3cd7099
2020-01-24 09:01:47 +00:00

111 lines
3.0 KiB
Go

// Copyright (C) 2020 Storj Labs, Inc.
// See LICENSE for copying information.
package cockroachutil
import (
"context"
"testing"
"github.com/stretchr/testify/require"
"github.com/zeebo/errs"
"storj.io/common/testcontext"
"storj.io/storj/private/dbutil/pgutil/pgtest"
"storj.io/storj/private/tagsql"
)
func TestLibPqCompatibility(t *testing.T) {
ctx := testcontext.New(t)
defer ctx.Cleanup()
if *pgtest.CrdbConnStr == "" {
t.Skip("CockroachDB flag missing")
}
testDB, err := OpenUnique(ctx, *pgtest.CrdbConnStr, "TestLibPqCompatibility")
require.NoError(t, err)
defer ctx.Check(testDB.Close)
// use a single dedicated conn for testing
conn, err := testDB.Conn(ctx)
require.NoError(t, err)
defer ctx.Check(conn.Close)
// should be in idle status, no transaction, initially
require.Equal(t, txnStatusIdle, getTxnStatus(ctx, t, conn))
require.False(t, checkIsInTx(ctx, t, conn))
// start a transaction
tx, err := conn.BeginTx(ctx, nil)
require.NoError(t, err)
func() {
defer func() { err = tx.Rollback() }()
// should be idle in transaction now
require.Equal(t, txnStatusIdleInTransaction, getTxnStatus(ctx, t, conn))
require.True(t, checkIsInTx(ctx, t, conn))
// issue successful query
rows, err := tx.QueryContext(ctx, `SELECT 1`)
require.NoError(t, err)
require.True(t, rows.Next())
var n int
err = rows.Scan(&n)
require.NoError(t, err)
require.False(t, rows.Next())
err = rows.Err()
require.NoError(t, err)
err = rows.Close()
require.NoError(t, err)
// should still be idle in transaction
require.Equal(t, txnStatusIdleInTransaction, getTxnStatus(ctx, t, conn))
require.True(t, checkIsInTx(ctx, t, conn))
// issue bad query
_, err = tx.QueryContext(ctx, `SELECT BALONEY SANDWICHES`)
require.Error(t, err)
// should be in a failed transaction now
require.Equal(t, txnStatusInFailedTransaction, getTxnStatus(ctx, t, conn))
require.True(t, checkIsInTx(ctx, t, conn))
}()
// check rollback error
require.NoError(t, err)
// should be back out of any transaction
require.Equal(t, txnStatusIdle, getTxnStatus(ctx, t, conn))
require.False(t, checkIsInTx(ctx, t, conn))
}
func withCockroachConn(ctx context.Context, sqlConn tagsql.Conn, fn func(conn *cockroachConn) error) error {
return sqlConn.Raw(ctx, func(rawConn interface{}) error {
crConn, ok := rawConn.(*cockroachConn)
if !ok {
return errs.New("conn object is %T, not *cockroachConn", crConn)
}
return fn(crConn)
})
}
func getTxnStatus(ctx context.Context, t *testing.T, sqlConn tagsql.Conn) (txnStatus transactionStatus) {
err := withCockroachConn(ctx, sqlConn, func(crConn *cockroachConn) error {
txnStatus = crConn.txnStatus()
return nil
})
require.NoError(t, err)
return txnStatus
}
func checkIsInTx(ctx context.Context, t *testing.T, sqlConn tagsql.Conn) (isInTx bool) {
err := withCockroachConn(ctx, sqlConn, func(crConn *cockroachConn) error {
isInTx = crConn.isInTransaction()
return nil
})
require.NoError(t, err)
return isInTx
}