netbird/management/internals/shared/db/conn.go

package db

import (
	"context"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	log "github.com/sirupsen/logrus"
	"gorm.io/gorm"
)

const (
	defaultTransactionTimeout = 5 * time.Minute
	connMaxLifetime           = time.Hour
	connMaxIdleTime           = 3 * time.Minute
)

// TxMetrics receives the duration of every committed top-level transaction.
type TxMetrics interface {
	CountTransactionDuration(duration time.Duration)
}

// Conn is the database connection shared by all repositories: one gorm handle,
// the pgx pool of a Postgres deployment and the engine they talk to.
type Conn struct {
	db        *gorm.DB
	pool      *pgxpool.Pool
	engine    Engine
	txTimeout time.Duration
	metrics   TxMetrics
}

// NewConn takes ownership of an open gorm handle and pool once it returns
// without error, applying the connection limits and transaction timeout
// configured through the environment.
func NewConn(ctx context.Context, gormDB *gorm.DB, engine Engine, pool *pgxpool.Pool) (*Conn, error) {
	sqlDB, err := gormDB.DB()
	if err != nil {
		return nil, err
	}

	txTimeout := defaultTransactionTimeout
	if v := os.Getenv("NB_STORE_TRANSACTION_TIMEOUT"); v != "" {
		if parsed, err := time.ParseDuration(v); err == nil {
			txTimeout = parsed
		}
	}
	log.WithContext(ctx).Infof("Setting transaction timeout to %v", txTimeout)

	conns := runtime.NumCPU()
	configuredConns, err := strconv.Atoi(os.Getenv("NB_SQL_MAX_OPEN_CONNS"))
	connsConfigured := err == nil
	if connsConfigured {
		conns = configuredConns
	}
	if engine == SqliteStoreEngine {
		if connsConfigured {
			log.WithContext(ctx).Warnf("setting NB_SQL_MAX_OPEN_CONNS is not supported for sqlite, using default value 1")
		}
		conns = 1
	}

	sqlDB.SetMaxOpenConns(conns)
	sqlDB.SetMaxIdleConns(conns)
	sqlDB.SetConnMaxLifetime(connMaxLifetime)
	sqlDB.SetConnMaxIdleTime(connMaxIdleTime)

	log.WithContext(ctx).Infof("Set max open db connections to %d, max idle to %d, max lifetime to %v, max idle time to %v",
		conns, conns, connMaxLifetime, connMaxIdleTime)

	return &Conn{db: gormDB, pool: pool, engine: engine, txTimeout: txTimeout}, nil
}

// DB returns the handle a query must run on: the transaction when tx is set,
// otherwise the shared connection.
func (c *Conn) DB(tx *Tx) *gorm.DB {
	if tx != nil {
		return tx.db
	}
	return c.db
}

// Pool returns the pgx pool for read paths that bypass gorm. It is nil on
// engines other than Postgres and inside a transaction, where the pool would
// not see the uncommitted writes.
func (c *Conn) Pool(tx *Tx) *pgxpool.Pool {
	if tx != nil {
		return nil
	}
	return c.pool
}

func (c *Conn) Engine() Engine {
	return c.engine
}

// SetTxMetrics registers the sink that receives transaction durations.
func (c *Conn) SetTxMetrics(metrics TxMetrics) {
	c.metrics = metrics
}

// AutoMigrate creates or updates the tables of the given models.
func (c *Conn) AutoMigrate(models ...any) error {
	return c.db.AutoMigrate(models...)
}

// Close releases the gorm connection and the pgx pool.
func (c *Conn) Close() error {
	if c.pool != nil {
		c.pool.Close()
	}
	sqlDB, err := c.db.DB()
	if err != nil {
		return fmt.Errorf("get db: %w", err)
	}
	return sqlDB.Close()
}