sql.go

package checkers

import (
	"context"
	"database/sql"
	"fmt"
)

//go:generate counterfeiter -o ../fakes/isqlpinger.go . SQLPinger
//go:generate counterfeiter -o ../fakes/isqlqueryer.go . SQLQueryer
//go:generate counterfeiter -o ../fakes/isqlexecer.go . SQLExecer

// SQLPinger is an interface that allows direct pinging of the database
type SQLPinger interface {
	PingContext(ctx context.Context) error
}

// SQLQueryer is an interface that allows querying of the database
type SQLQueryer interface {
	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
}

// SQLExecer is an interface that allows executing of queries in the database
type SQLExecer interface {
	ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
}

// SQLQueryerResultHandler is the BYO function to
// handle the result of an SQL SELECT query
type SQLQueryerResultHandler func(rows *sql.Rows) (bool, error)

// SQLExecerResultHandler is the BYO function
// to handle a database exec result
type SQLExecerResultHandler func(result sql.Result) (bool, error)

// SQLConfig is used for configuring a database check.
// One of the Pinger, Queryer, or Execer fields is required.
//
//
// If Execer is set, it will take precedence over Queryer and Pinger,
// Execer implements the SQLExecer interface in this package.
// The sql.DB and sql.TX structs both implement this interface.
//
// Note that if the Execer is set, then the ExecerResultHandler
// and Query values MUST also be set
//
//
// If Queryer is set, it will take precedence over Pinger.
// SQLQueryer implements the SQLQueryer interface in this package.
// The sql.DB and sql.TX structs both implement this interface.
//
// Note that if the Queryer is set, then the QueryerResultHandler
// and Query values MUST also be set
//
//
// Pinger implements the SQLPinger interface in this package.
// The sql.DB struct implements this interface.
type SQLConfig struct {
	// Pinger is the value implementing SQLPinger
	Pinger SQLPinger

	// Queryer is the value implementing SQLQueryer
	Queryer SQLQueryer

	// Execer is the value implementing SQLExecer
	Execer SQLExecer

	// Query is the parameterized SQL query required
	// with both Queryer and Execer
	Query string

	// Params are the SQL query parameters, if any
	Params []interface{}

	// QueryerResultHandler handles the result of
	// the QueryContext function
	QueryerResultHandler SQLQueryerResultHandler

	// ExecerResultHandler handles the result of
	// the ExecContext function
	ExecerResultHandler SQLExecerResultHandler
}

// SQL implements the "ICheckable" interface
type SQL struct {
	Config *SQLConfig
}

// NewSQL creates a new database checker that can be used for ".AddCheck(s)".
func NewSQL(cfg *SQLConfig) (*SQL, error) {
	if err := validateSQLConfig(cfg); err != nil {
		return nil, err
	}

	return &SQL{
		Config: cfg,
	}, nil
}

// DefaultQueryHandler is the default SQLQueryer result handler
// that assumes one row was returned from the passed query
func DefaultQueryHandler(rows *sql.Rows) (bool, error) {
	defer rows.Close()

	numRows := 0
	for rows.Next() {
		numRows++
	}

	return numRows == 1, nil
}

// DefaultExecHandler is the default SQLExecer result handler
// that assumes one row was affected in the passed query
func DefaultExecHandler(result sql.Result) (bool, error) {
	affectedRows, err := result.RowsAffected()
	if err != nil {
		return false, err
	}

	return affectedRows == int64(1), nil
}

// this makes sure the sql check is properly configured
func validateSQLConfig(cfg *SQLConfig) error {
	if cfg == nil {
		return fmt.Errorf("config is required")
	}

	if cfg.Execer == nil && cfg.Queryer == nil && cfg.Pinger == nil {
		return fmt.Errorf("one of Execer, Queryer, or Pinger is required in SQLConfig")
	}

	if (cfg.Execer != nil || cfg.Queryer != nil) && len(cfg.Query) == 0 {
		return fmt.Errorf("SQLConfig.Query is required")
	}

	return nil
}

// Status is used for performing a database ping against a dependency; it satisfies
// the "ICheckable" interface.
func (s *SQL) Status() (interface{}, error) {
	if err := validateSQLConfig(s.Config); err != nil {
		return nil, err
	}

	switch {
	// check for SQLExecer first
	case s.Config.Execer != nil:
		// if the result handler is nil, use the default
		if s.Config.ExecerResultHandler == nil {
			s.Config.ExecerResultHandler = DefaultExecHandler
		}
		// run the execer
		return s.runExecer()
	// check for SQLQueryer next
	case s.Config.Queryer != nil:
		// if the result handler is nil, use the default
		if s.Config.QueryerResultHandler == nil {
			s.Config.QueryerResultHandler = DefaultQueryHandler
		}
		// run the queryer
		return s.runQueryer()
	// finally, must be a pinger
	default:
		ctx := context.Background()
		return nil, s.Config.Pinger.PingContext(ctx)
	}
}

// This will run the execer from the Status func
func (s *SQL) runExecer() (interface{}, error) {
	ctx := context.Background()
	result, err := s.Config.Execer.ExecContext(ctx, s.Config.Query, s.Config.Params...)
	if err != nil {
		return nil, err
	}

	ok, err := s.Config.ExecerResultHandler(result)
	if err != nil {
		return nil, err
	}

	if !ok {
		return nil, fmt.Errorf("userland exec result handler returned false")
	}

	return nil, nil
}

// This will run the queryer from the Status func
func (s *SQL) runQueryer() (interface{}, error) {
	ctx := context.Background()
	rows, err := s.Config.Queryer.QueryContext(ctx, s.Config.Query, s.Config.Params...)
	if err != nil {
		return nil, err
	}

	// the BYO result handler is responsible for closing the rows

	ok, err := s.Config.QueryerResultHandler(rows)
	if err != nil {
		return nil, err
	}

	if !ok {
		return nil, fmt.Errorf("userland query result handler returned false")
	}

	return nil, nil
}