code.vegaprotocol.io/vega@v0.79.0/libs/config/dbconnection.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 config
    17  
    18  import (
    19  	"fmt"
    20  	"strconv"
    21  
    22  	"code.vegaprotocol.io/vega/libs/config/encoding"
    23  
    24  	"github.com/jackc/pgx/v4/pgxpool"
    25  )
    26  
    27  type PostgresConnection struct {
    28  	Host             string            `description:"hostname of postgres server"                                               long:"host"`
    29  	Port             int               `description:"port postgres is running on"                                               long:"port"`
    30  	Username         string            `description:"username to connect with"                                                  long:"username"`
    31  	Password         string            `description:"password for user"                                                         long:"password"`
    32  	Database         string            `description:"database name"                                                             long:"database"`
    33  	SocketDir        string            `description:"location of postgres UNIX socket directory (used if host is empty string)" long:"socket-dir"`
    34  	ApplicationName  string            `description:"identify the application to the database using this name"                  long:"application-name"`
    35  	StatementTimeout encoding.Duration `description:"Terminate any database connections that take longer than this"             long:"statement-timeout"`
    36  }
    37  
    38  func (conf PostgresConnection) ToConnectionString() string {
    39  	if conf.Host == "" {
    40  		//nolint:nosprintfhostport
    41  		return fmt.Sprintf("postgresql://%s:%s@/%s?host=%s&port=%d",
    42  			conf.Username,
    43  			conf.Password,
    44  			conf.Database,
    45  			conf.SocketDir,
    46  			conf.Port)
    47  	}
    48  	//nolint:nosprintfhostport
    49  	return fmt.Sprintf("postgresql://%s:%s@%s:%d/%s",
    50  		conf.Username,
    51  		conf.Password,
    52  		conf.Host,
    53  		conf.Port,
    54  		conf.Database)
    55  }
    56  
    57  func (conf PostgresConnection) ToPgxPoolConfig() (*pgxpool.Config, error) {
    58  	cfg, err := pgxpool.ParseConfig(conf.ToConnectionString())
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  
    63  	if conf.ApplicationName != "" {
    64  		cfg.ConnConfig.RuntimeParams["application_name"] = "Vega Data Node"
    65  	}
    66  
    67  	if conf.StatementTimeout.Get() > 0 {
    68  		cfg.ConnConfig.RuntimeParams["statement_timeout"] = strconv.Itoa(int(conf.StatementTimeout.Get().Milliseconds()))
    69  	}
    70  
    71  	return cfg, nil
    72  }