github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/lorry/engines/postgres/config.go (about)

     1  /*
     2  Copyright (C) 2022-2023 ApeCloud Co., Ltd
     3  
     4  This file is part of KubeBlocks project
     5  
     6  This program is free software: you can redistribute it and/or modify
     7  it under the terms of the GNU Affero General Public License as published by
     8  the Free Software Foundation, either version 3 of the License, or
     9  (at your option) any later version.
    10  
    11  This program is distributed in the hope that it will be useful
    12  but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  GNU Affero General Public License for more details.
    15  
    16  You should have received a copy of the GNU Affero General Public License
    17  along with this program.  If not, see <http://www.gnu.org/licenses/>.
    18  */
    19  
    20  package postgres
    21  
    22  import (
    23  	"fmt"
    24  
    25  	"github.com/jackc/pgx/v5/pgxpool"
    26  	"github.com/pkg/errors"
    27  
    28  	"github.com/1aal/kubeblocks/pkg/lorry/dcs"
    29  	viper "github.com/1aal/kubeblocks/pkg/viperx"
    30  )
    31  
    32  const (
    33  	ConnectionURLKey = "url"
    34  	DefaultPort      = 5432
    35  )
    36  
    37  type Config struct {
    38  	URL            string
    39  	Username       string
    40  	Password       string
    41  	Host           string
    42  	Port           int
    43  	Database       string
    44  	MaxConnections int32
    45  	MinConnections int32
    46  	pgxConfig      *pgxpool.Config
    47  }
    48  
    49  var config *Config
    50  
    51  func NewConfig(properties map[string]string) (*Config, error) {
    52  	config = &Config{}
    53  
    54  	url, ok := properties[ConnectionURLKey]
    55  	if !ok || url == "" {
    56  		return nil, errors.Errorf("required metadata not set: %s", ConnectionURLKey)
    57  	}
    58  
    59  	poolConfig, err := pgxpool.ParseConfig(url)
    60  	if err != nil {
    61  		return nil, errors.Errorf("error opening DB connection: %v", err)
    62  	}
    63  
    64  	config.Username = poolConfig.ConnConfig.User
    65  	config.Password = poolConfig.ConnConfig.Password
    66  	config.Host = poolConfig.ConnConfig.Host
    67  	config.Port = int(poolConfig.ConnConfig.Port)
    68  	config.Database = poolConfig.ConnConfig.Database
    69  	config.MaxConnections = poolConfig.MaxConns
    70  	config.MinConnections = poolConfig.MinConns
    71  
    72  	if viper.IsSet("KB_SERVICE_USER") {
    73  		config.Username = viper.GetString("KB_SERVICE_USER")
    74  	}
    75  	if viper.IsSet("KB_SERVICE_PASSWORD") {
    76  		config.Password = viper.GetString("KB_SERVICE_PASSWORD")
    77  	}
    78  
    79  	config.URL = config.GetConnectURLWithHost(config.Host)
    80  	pgxConfig, _ := pgxpool.ParseConfig(config.URL)
    81  	config.pgxConfig = pgxConfig
    82  
    83  	return config, nil
    84  }
    85  
    86  func (config *Config) GetDBPort() int {
    87  	if config.Port == 0 {
    88  		return DefaultPort
    89  	}
    90  
    91  	return config.Port
    92  }
    93  
    94  func (config *Config) GetConnectURLWithHost(host string) string {
    95  	return fmt.Sprintf("user=%s password=%s host=%s port=%d dbname=%s",
    96  		config.Username, config.Password, host, config.Port, config.Database)
    97  }
    98  
    99  func (config *Config) GetConsensusIPPort(cluster *dcs.Cluster, name string) string {
   100  	return fmt.Sprintf("%s.%s-headless.%s.svc:1%d", name, cluster.ClusterCompName, cluster.Namespace, config.GetDBPort())
   101  }