github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/lorry/engines/wesql/conn.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 wesql
    21  
    22  import (
    23  	"context"
    24  	"database/sql"
    25  	"strings"
    26  
    27  	"github.com/pkg/errors"
    28  
    29  	"github.com/1aal/kubeblocks/pkg/lorry/dcs"
    30  )
    31  
    32  // GetDBConnWithMember retrieves a database connection for a specific member of a cluster.
    33  func (mgr *Manager) GetDBConnWithMember(cluster *dcs.Cluster, member *dcs.Member) (*sql.DB, error) {
    34  	var db *sql.DB
    35  	var err error
    36  	if member != nil {
    37  		addr := cluster.GetMemberAddrWithPort(*member)
    38  		db, err = config.GetDBConnWithAddr(addr)
    39  		if err != nil {
    40  			return nil, errors.Wrap(err, "new db connection failed")
    41  		}
    42  	}
    43  	return db, nil
    44  }
    45  
    46  // GetLeaderConn retrieves a database connection to the leader member of a cluster.
    47  func (mgr *Manager) GetLeaderConn(ctx context.Context, cluster *dcs.Cluster) (*sql.DB, error) {
    48  	mgr.Logger.Info("Get leaader from dcs cluster")
    49  	leaderMember := cluster.GetLeaderMember()
    50  	if leaderMember == nil {
    51  		mgr.Logger.Info("Get leaader from db cluster local")
    52  		leaderMember = mgr.GetLeaderMember(ctx, cluster)
    53  	}
    54  	if leaderMember == nil {
    55  		return nil, errors.New("the cluster has no leader")
    56  	}
    57  	return mgr.GetDBConnWithMember(cluster, leaderMember)
    58  }
    59  
    60  // GetLeaderMember retrieves the leader member of a cluster
    61  func (mgr *Manager) GetLeaderMember(ctx context.Context, cluster *dcs.Cluster) *dcs.Member {
    62  	clusterLocalInfo, err := mgr.GetClusterLocalInfo(ctx)
    63  	if err != nil || clusterLocalInfo == nil {
    64  		mgr.Logger.Error(err, "Get cluster local info failed")
    65  		return nil
    66  	}
    67  
    68  	leaderAddr := clusterLocalInfo.GetString("CURRENT_LEADER")
    69  	if leaderAddr == "" {
    70  		return nil
    71  	}
    72  	leaderParts := strings.Split(leaderAddr, ".")
    73  	if len(leaderParts) > 0 {
    74  		return cluster.GetMemberWithName(leaderParts[0])
    75  	}
    76  
    77  	return nil
    78  }