storj/private/dbutil/cockroachutil/driver.go
Egon Elbre 4e8d53c8fb private/dbutil/pgutil: ensure storagenode doesn't depend on pgx
pgx is a large dependency and there's no need to include it in
storagenode binary.

Change-Id: I49c304c6420733d5f095d7edb35d32811210e41a
2020-09-30 14:28:47 +00:00

348 lines
10 KiB
Go

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.
package cockroachutil
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"io"
"strings"
"github.com/jackc/pgx/v4/stdlib"
"github.com/zeebo/errs"
"storj.io/storj/private/dbutil/pgutil/pgerrcode"
)
// Driver is the type for the "cockroach" sql/database driver. It uses
// github.com/jackc/pgx/v4/stdlib under the covers because of Cockroach's
// PostgreSQL compatibility, but allows differentiation between pg and crdb
// connections.
type Driver struct {
pgxDriver stdlib.Driver
}
// Open opens a new cockroachDB connection.
func (cd *Driver) Open(name string) (driver.Conn, error) {
name = translateName(name)
conn, err := cd.pgxDriver.Open(name)
if err != nil {
return nil, err
}
pgxStdlibConn, ok := conn.(*stdlib.Conn)
if !ok {
return nil, errs.New("Conn from pgx is not a *stdlib.Conn??? T: %T", conn)
}
return &cockroachConn{pgxStdlibConn}, nil
}
// OpenConnector obtains a new db Connector, which sql.DB can use to
// obtain each needed connection at the appropriate time.
func (cd *Driver) OpenConnector(name string) (driver.Connector, error) {
name = translateName(name)
pgxConnector, err := cd.pgxDriver.OpenConnector(name)
if err != nil {
return nil, err
}
return &cockroachConnector{driver: cd, pgxConnector: pgxConnector}, nil
}
// cockroachConnector is a thin wrapper around a pq-based connector. This allows
// Driver to supply our custom cockroachConn type for connections.
type cockroachConnector struct {
driver *Driver
pgxConnector driver.Connector
}
// Driver returns the driver being used for this connector.
func (c *cockroachConnector) Driver() driver.Driver {
return c.driver
}
// Connect creates a new connection using the connector.
func (c *cockroachConnector) Connect(ctx context.Context) (driver.Conn, error) {
pgxConn, err := c.pgxConnector.Connect(ctx)
if err != nil {
return nil, err
}
pgxStdlibConn, ok := pgxConn.(*stdlib.Conn)
if !ok {
return nil, errs.New("Conn from pgx is not a *stdlib.Conn??? T: %T", pgxConn)
}
return &cockroachConn{pgxStdlibConn}, nil
}
type connAll interface {
driver.Conn
driver.ConnBeginTx
driver.ExecerContext
driver.QueryerContext
}
// cockroachConn is a connection to a database. It is not used concurrently by multiple goroutines.
type cockroachConn struct {
underlying *stdlib.Conn
}
// Assert that cockroachConn fulfills connAll.
var _ connAll = (*cockroachConn)(nil)
// Close closes the cockroachConn.
func (c *cockroachConn) Close() error {
return c.underlying.Close()
}
// ExecContext (when implemented by a driver.Conn) provides ExecContext
// functionality to a sql.DB instance. This implementation provides
// retry semantics for single statements.
func (c *cockroachConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
result, err := c.underlying.ExecContext(ctx, query, args)
for err != nil && !c.isInTransaction() && NeedsRetry(err) {
mon.Event("needed_retry")
result, err = c.underlying.ExecContext(ctx, query, args)
}
return result, err
}
type cockroachRows struct {
rows driver.Rows
firstResults []driver.Value
eof bool
}
// Columns returns the names of the columns.
func (rows *cockroachRows) Columns() []string {
return rows.rows.Columns()
}
// Close closes the rows iterator.
func (rows *cockroachRows) Close() error {
return rows.rows.Close()
}
// Next implements the Next method on driver.Rows.
func (rows *cockroachRows) Next(dest []driver.Value) error {
if rows.eof {
return io.EOF
}
if rows.firstResults == nil {
return rows.rows.Next(dest)
}
copy(dest, rows.firstResults)
rows.firstResults = nil
return nil
}
func wrapRows(rows driver.Rows) (crdbRows *cockroachRows, err error) {
columns := rows.Columns()
dest := make([]driver.Value, len(columns))
err = rows.Next(dest)
if err != nil {
if errors.Is(err, io.EOF) {
return &cockroachRows{rows: rows, firstResults: nil, eof: true}, nil
}
return nil, err
}
return &cockroachRows{rows: rows, firstResults: dest}, nil
}
// QueryContext (when implemented by a driver.Conn) provides QueryContext
// functionality to a sql.DB instance. This implementation provides
// retry semantics for single statements.
func (c *cockroachConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (_ driver.Rows, err error) {
defer mon.Task()(&ctx)(&err)
for {
result, err := c.underlying.QueryContext(ctx, query, args)
if err != nil {
if NeedsRetry(err) {
if c.isInTransaction() {
return nil, err
}
mon.Event("needed_retry")
continue
}
return nil, err
}
wrappedResult, err := wrapRows(result)
if err != nil {
// If this returns an error it's probably the same error
// we got from calling Next inside wrapRows.
_ = result.Close()
if NeedsRetry(err) {
if c.isInTransaction() {
return nil, err
}
mon.Event("needed_retry")
continue
}
return nil, err
}
return wrappedResult, nil
}
}
// Begin starts a new transaction.
func (c *cockroachConn) Begin() (driver.Tx, error) {
return c.BeginTx(context.Background(), driver.TxOptions{})
}
// BeginTx begins a new transaction using the specified context and with the specified options.
func (c *cockroachConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
return c.underlying.BeginTx(ctx, opts)
}
// Prepare prepares a statement for future execution.
func (c *cockroachConn) Prepare(query string) (driver.Stmt, error) {
pqStmt, err := c.underlying.Prepare(query)
if err != nil {
return nil, err
}
adapted, ok := pqStmt.(stmtAll)
if !ok {
return nil, errs.New("Stmt type %T does not provide stmtAll?!", adapted)
}
return &cockroachStmt{underlyingStmt: adapted, conn: c}, nil
}
type transactionStatus byte
const (
txnStatusIdle transactionStatus = 'I'
txnStatusIdleInTransaction transactionStatus = 'T'
txnStatusInFailedTransaction transactionStatus = 'E'
)
func (c *cockroachConn) txnStatus() transactionStatus {
pgConn := c.underlying.Conn().PgConn()
return transactionStatus(pgConn.TxStatus())
}
func (c *cockroachConn) isInTransaction() bool {
txnStatus := c.txnStatus()
return txnStatus == txnStatusIdleInTransaction || txnStatus == txnStatusInFailedTransaction
}
type stmtAll interface {
driver.Stmt
driver.StmtExecContext
driver.StmtQueryContext
}
type cockroachStmt struct {
underlyingStmt stmtAll
conn *cockroachConn
}
// Assert that cockroachStmt satisfies StmtExecContext and StmtQueryContext.
var _ stmtAll = (*cockroachStmt)(nil)
// Close closes a prepared statement.
func (stmt *cockroachStmt) Close() error {
return stmt.underlyingStmt.Close()
}
// NumInput returns the number of placeholder parameters.
func (stmt *cockroachStmt) NumInput() int {
return stmt.underlyingStmt.NumInput()
}
// Exec executes a SQL statement in the background context.
func (stmt *cockroachStmt) Exec(args []driver.Value) (driver.Result, error) {
// since (driver.Stmt).Exec() is deprecated, we translate our Value args to NamedValue args
// and pass in background context to ExecContext instead.
namedArgs := make([]driver.NamedValue, len(args))
for i, arg := range args {
namedArgs[i] = driver.NamedValue{Ordinal: i + 1, Value: arg}
}
result, err := stmt.underlyingStmt.ExecContext(context.Background(), namedArgs)
for err != nil && !stmt.conn.isInTransaction() && NeedsRetry(err) {
mon.Event("needed_retry")
result, err = stmt.underlyingStmt.ExecContext(context.Background(), namedArgs)
}
return result, err
}
// Query executes a query in the background context.
func (stmt *cockroachStmt) Query(args []driver.Value) (driver.Rows, error) {
// since (driver.Stmt).Query() is deprecated, we translate our Value args to NamedValue args
// and pass in background context to QueryContext instead.
namedArgs := make([]driver.NamedValue, len(args))
for i, arg := range args {
namedArgs[i] = driver.NamedValue{Ordinal: i + 1, Value: arg}
}
return stmt.QueryContext(context.Background(), namedArgs)
}
// ExecContext executes SQL statements in the specified context.
func (stmt *cockroachStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
result, err := stmt.underlyingStmt.ExecContext(ctx, args)
for err != nil && !stmt.conn.isInTransaction() && NeedsRetry(err) {
mon.Event("needed_retry")
result, err = stmt.underlyingStmt.ExecContext(ctx, args)
}
return result, err
}
// QueryContext executes a query in the specified context.
func (stmt *cockroachStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (_ driver.Rows, err error) {
defer mon.Task()(&ctx)(&err)
for {
result, err := stmt.underlyingStmt.QueryContext(ctx, args)
if err != nil {
if NeedsRetry(err) {
if stmt.conn.isInTransaction() {
return nil, err
}
mon.Event("needed_retry")
continue
}
return nil, err
}
wrappedResult, err := wrapRows(result)
if err != nil {
// If this returns an error it's probably the same error
// we got from calling Next inside wrapRows.
_ = result.Close()
if NeedsRetry(err) {
if stmt.conn.isInTransaction() {
return nil, err
}
mon.Event("needed_retry")
continue
}
return nil, err
}
return wrappedResult, nil
}
}
// translateName changes the scheme name in a `cockroach://` URL to
// `postgres://`, as that is what jackc/pgx will expect.
func translateName(name string) string {
if strings.HasPrefix(name, "cockroach://") {
name = "postgres://" + name[12:]
}
return name
}
// NeedsRetry checks if the error code means a retry is needed,
// borrowed from code in crdb.
func NeedsRetry(err error) bool {
code := pgerrcode.FromError(err)
// 57P01 occurs when a CRDB node rejoins the cluster but is not ready to accept connections
// CRDB support recommended a retry at this point
// Support ticket: https://support.cockroachlabs.com/hc/en-us/requests/5510
// TODO re-evaluate this if support provides a better solution
return code == "40001" || code == "CR000" || code == "57P01"
}
var defaultDriver = &Driver{}
func init() {
sql.Register("cockroach", defaultDriver)
}