threshold/internal/db/db.go

package db

import (
	"context"
	"database/sql"
	stderrors "errors"
	"fmt"

	"github.com/hashicorp/boundary/internal/errors"
	"github.com/hashicorp/boundary/internal/observability/event"
	"github.com/hashicorp/go-hclog"
	"github.com/jackc/pgconn"
	_ "github.com/jackc/pgx/v4"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

type DbType int

const (
	UnknownDB DbType = 0
	Postgres  DbType = 1
)

func (db DbType) String() string {
	return [...]string{
		"unknown",
		"postgres",
	}[db]
}
func StringToDbType(dialect string) (DbType, error) {
	switch dialect {
	case "postgres":
		return Postgres, nil
	default:
		return UnknownDB, fmt.Errorf("%s is an unknown dialect", dialect)
	}
}

// DB is a wrapper around the ORM
type DB struct {
	*gorm.DB
}

// SqlDB returns the underlying sql.DB
//
// Note: This func is not named DB(), because our choice of ORM (gorm) which is
// embedded for various reasons, already exports a DB() func and it's base type
// is also an export DB type
func (d *DB) SqlDB(ctx context.Context) (*sql.DB, error) {
	const op = "db.(DB).SqlDB"
	if d.DB == nil {
		return nil, errors.New(ctx, errors.Internal, op, "missing underlying database")
	}
	return d.DB.DB()
}

// Close the underlying sql.DB
func (d *DB) Close(ctx context.Context) error {
	const op = "db.(DB).Close"
	if d.DB == nil {
		return errors.New(ctx, errors.Internal, op, "missing underlying database")
	}
	underlying, err := d.DB.DB()
	if err != nil {
		return errors.Wrap(ctx, err, op)
	}
	return underlying.Close()
}

// Open a database connection which is long-lived. The options of
// WithGormFormatter and WithMaxOpenConnections are supported.
//
// Note: Consider if you need to call Close() on the returned DB.  Typically the
// answer is no, but there are occasions when it's necessary.  See the sql.DB
// docs for more information.
func Open(dbType DbType, connectionUrl string, opt ...Option) (*DB, error) {
	var dialect gorm.Dialector
	switch dbType {
	case Postgres:
		dialect = postgres.New(postgres.Config{
			DSN: connectionUrl},
		)
	default:
		return nil, fmt.Errorf("unable to open %s database type", dbType)
	}
	db, err := gorm.Open(dialect, &gorm.Config{
		ConvertNullToZeroValues: true,
	})
	if err != nil {
		return nil, fmt.Errorf("unable to open database: %w", err)
	}
	opts := GetOpts(opt...)
	if opts.withGormFormatter != nil {
		newLogger := logger.New(
			getGormLogger(opts.withGormFormatter),
			logger.Config{
				LogLevel: logger.Error, // Log level
				Colorful: false,        // Disable color
			},
		)
		db = db.Session(&gorm.Session{Logger: newLogger})
	}
	if opts.withMaxOpenConnections > 0 {
		if dialect.Name() == "postgres" && opts.withMaxOpenConnections == 1 {
			return nil, fmt.Errorf("unable to create db object with dialect %s: %s", dialect, "max_open_connections must be unlimited by setting 0 or at least 2")
		}
		underlyingDB, err := db.DB()
		if err != nil {
			return nil, fmt.Errorf("unable retreive db: %w", err)
		}
		underlyingDB.SetMaxOpenConns(opts.withMaxOpenConnections)
	}
	return &DB{db}, nil
}

func GetGormLogFormatter(log hclog.Logger) func(values ...interface{}) (messages []interface{}) {
	const op = "db.GetGormLogFormatter"
	ctx := context.TODO()
	return func(values ...interface{}) (messages []interface{}) {
		if len(values) > 2 && values[0].(string) == "log" {
			switch values[2].(type) {
			case *pgconn.PgError:
				if log.IsTrace() {
					event.WriteError(ctx, op, stderrors.New("error from database adapter"), event.WithInfo("error", values[2], "location", values[1]))
				}
			}
			return nil
		}
		return nil
	}
}

type gormLogger struct {
	logger hclog.Logger
}

func (g gormLogger) Printf(msg string, values ...interface{}) {
	if len(values) > 1 {
		switch values[1].(type) {
		case *pgconn.PgError:
			g.logger.Trace("error from database adapter", "location", values[0], "error", values[1])
		}
	}
}

func getGormLogger(log hclog.Logger) gormLogger {
	return gormLogger{logger: log}
}