code.vegaprotocol.io/vega@v0.79.0/datanode/sqlstore/connection_source.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package sqlstore
    17  
    18  import (
    19  	"context"
    20  	"fmt"
    21  	"strconv"
    22  
    23  	"code.vegaprotocol.io/vega/datanode/entities"
    24  	"code.vegaprotocol.io/vega/libs/num"
    25  
    26  	"github.com/jackc/pgconn"
    27  	"github.com/jackc/pgtype"
    28  	shopspring "github.com/jackc/pgtype/ext/shopspring-numeric"
    29  	"github.com/jackc/pgx/v4"
    30  	"github.com/jackc/pgx/v4/pgxpool"
    31  	"github.com/pkg/errors"
    32  )
    33  
    34  var (
    35  	numSpareConnections = 15 // If possible, the pool size will be (max_connections - numSpareConnections).
    36  	poolSizeLowerBound  = 10 // But it will never be lower than this.
    37  )
    38  
    39  type Connection interface {
    40  	Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error)
    41  	QueryRow(ctx context.Context, sql string, args ...interface{}) pgx.Row
    42  	QueryFunc(ctx context.Context, sql string, args []interface{}, scans []interface{}, f func(pgx.QueryFuncRow) error) (pgconn.CommandTag, error)
    43  	SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults
    44  	CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error)
    45  	Exec(ctx context.Context, sql string, arguments ...interface{}) (pgconn.CommandTag, error)
    46  }
    47  
    48  func setMaxPoolSize(ctx context.Context, poolConfig *pgxpool.Config, conf ConnectionConfig) error {
    49  	conn, err := pgx.Connect(ctx, poolConfig.ConnString())
    50  	if err != nil {
    51  		return fmt.Errorf("connecting to db: %w", err)
    52  	}
    53  	defer conn.Close(ctx)
    54  
    55  	var maxConnectionsStr string
    56  	if err := conn.QueryRow(ctx, "SHOW max_connections;").Scan(&maxConnectionsStr); err != nil {
    57  		return fmt.Errorf("querying max_connections: %w", err)
    58  	}
    59  
    60  	maxConnections, err := strconv.Atoi(maxConnectionsStr)
    61  	if err != nil {
    62  		return fmt.Errorf("max_connections was not an integer: %w", err)
    63  	}
    64  
    65  	maxConnections = num.MaxV(maxConnections-numSpareConnections, poolSizeLowerBound)
    66  	if conf.MaxConnPoolSize > 0 && maxConnections > conf.MaxConnPoolSize {
    67  		maxConnections = conf.MaxConnPoolSize
    68  	}
    69  
    70  	poolConfig.MaxConns = int32(maxConnections)
    71  	return nil
    72  }
    73  
    74  func wrapE(err error) error {
    75  	switch {
    76  	case errors.Is(err, pgx.ErrNoRows):
    77  		return entities.ErrNotFound
    78  	case errors.Is(err, entities.ErrInvalidID):
    79  		return entities.ErrInvalidID
    80  	default:
    81  		return err
    82  	}
    83  }
    84  
    85  func registerNumericType(poolConfig *pgxpool.Config) {
    86  	// Cause postgres numeric types to be loaded as shopspring decimals and vice-versa
    87  	poolConfig.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
    88  		conn.ConnInfo().RegisterDataType(pgtype.DataType{
    89  			Value: &shopspring.Numeric{},
    90  			Name:  "numeric",
    91  			OID:   pgtype.NumericOID,
    92  		})
    93  		return nil
    94  	}
    95  }
    96  
    97  func CreateConnectionPool(ctx context.Context, conf ConnectionConfig) (*pgxpool.Pool, error) {
    98  	poolConfig, err := conf.GetPoolConfig()
    99  	if err != nil {
   100  		return nil, fmt.Errorf("failed to get pool config: %w", err)
   101  	}
   102  
   103  	setMaxPoolSize(ctx, poolConfig, conf)
   104  	registerNumericType(poolConfig)
   105  
   106  	poolConfig.MinConns = conf.MinConnPoolSize
   107  
   108  	pool, err := pgxpool.ConnectConfig(ctx, poolConfig)
   109  	if err != nil {
   110  		return nil, fmt.Errorf("error connecting to database: %w", err)
   111  	}
   112  
   113  	return pool, nil
   114  }