github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/conf/postgres/postgres.go (about)

     1  package postgres
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/pkg/errors"
     9  
    10  	"github.com/machinefi/w3bstream/pkg/depends/base/types"
    11  	"github.com/machinefi/w3bstream/pkg/depends/kit/sqlx"
    12  	"github.com/machinefi/w3bstream/pkg/depends/kit/sqlx/driver/postgres"
    13  	"github.com/machinefi/w3bstream/pkg/depends/x/misc/retry"
    14  )
    15  
    16  type Endpoint struct {
    17  	Master   types.Endpoint
    18  	Slave    types.Endpoint
    19  	Database *sqlx.Database `env:"-"`
    20  	Retry    *retry.Retry
    21  
    22  	Extensions      []string
    23  	PoolSize        int
    24  	ConnMaxLifetime types.Duration
    25  
    26  	*sqlx.DB `env:"-"`
    27  	slave    *sqlx.DB `env:"-"`
    28  }
    29  
    30  var _ types.DefaultSetter = (*Endpoint)(nil)
    31  
    32  func (e Endpoint) LivenessCheck() map[string]string {
    33  	s := map[string]string{}
    34  
    35  	_, err := e.DB.ExecContext(context.Background(), "SELECT 1")
    36  	if err != nil {
    37  		s[e.Master.Host()] = err.Error()
    38  	} else {
    39  		s[e.Master.Host()] = "ok"
    40  	}
    41  
    42  	if e.slave != nil {
    43  		_, err := e.slave.ExecContext(context.Background(), "SELECT 1")
    44  		if err != nil {
    45  			s[e.Slave.Host()] = err.Error()
    46  		} else {
    47  			s[e.Slave.Host()] = "ok"
    48  		}
    49  	}
    50  
    51  	return s
    52  }
    53  
    54  func (e *Endpoint) SetDefault() {
    55  	if e.PoolSize == 0 {
    56  		e.PoolSize = 10
    57  	}
    58  	if e.ConnMaxLifetime == 0 {
    59  		e.ConnMaxLifetime = types.Duration(time.Hour)
    60  	}
    61  	if e.Master.IsZero() {
    62  		e.Master.Hostname, e.Master.Port = "127.0.0.1", 5432
    63  	}
    64  	e.Master.Scheme = "postgres"
    65  	if e.Database.Name == "" && len(e.Master.Base) > 0 {
    66  		e.Database.Name = e.Master.Base
    67  	}
    68  	if e.Retry == nil {
    69  		e.Retry = retry.Default
    70  	}
    71  }
    72  
    73  func (e Endpoint) UseSlave() sqlx.DBExecutor {
    74  	if e.slave == nil {
    75  		return e.DB
    76  	}
    77  	return e.slave
    78  }
    79  
    80  func (e *Endpoint) conn(url string, readonly bool) (*sqlx.DB, error) {
    81  	connector := &postgres.Connector{
    82  		Host:  url,
    83  		Extra: e.Master.Param.Encode(),
    84  	}
    85  	if e.Database != nil {
    86  		connector.MustSchema = e.Database.Schema
    87  	}
    88  	if !readonly {
    89  		connector.Extensions = e.Extensions
    90  	}
    91  	db := e.Database.OpenDB(connector)
    92  	db.SetMaxOpenConns(e.PoolSize)
    93  	db.SetMaxIdleConns(e.PoolSize / 2)
    94  	db.SetConnMaxLifetime(e.ConnMaxLifetime.Duration())
    95  
    96  	_, err := db.ExecContext(context.Background(), "SELECT 1")
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  	return db, nil
   101  }
   102  
   103  func (e *Endpoint) Init() error {
   104  	// cover default database name
   105  	if len(e.Master.Base) > 0 {
   106  		e.Database.Name = e.Master.Base
   107  	}
   108  	// try master
   109  	err := e.Retry.Do(func() error {
   110  		db, err := e.conn(e.masterURL(), false)
   111  		if err != nil {
   112  			return err
   113  		}
   114  		e.DB = db
   115  		return nil
   116  	})
   117  	if err != nil {
   118  		return errors.Errorf(
   119  			"try to connect master: %s failed: %v",
   120  			e.masterURL(), err,
   121  		)
   122  	}
   123  	// try slave if config
   124  	if !e.Slave.IsZero() {
   125  		if err = e.Retry.Do(func() error {
   126  			db, err := e.conn(e.slaveURL(), false)
   127  			if err != nil {
   128  				return err
   129  			}
   130  			e.slave = db
   131  			return nil
   132  		}); err != nil {
   133  			return errors.Errorf(
   134  				"try to connect master: %s failed: %v",
   135  				e.masterURL(), err,
   136  			)
   137  		}
   138  	}
   139  	return nil
   140  }
   141  
   142  func (e Endpoint) masterURL() string {
   143  	passwd := e.Master.Password
   144  	if passwd != "" {
   145  		passwd = ":" + passwd
   146  	}
   147  	return fmt.Sprintf("postgres://%s%s@%s", e.Master.Username, passwd, e.Master.Host())
   148  }
   149  
   150  func (e Endpoint) slaveURL() string {
   151  	passwd := e.Master.Password
   152  	if passwd != "" {
   153  		passwd = ":" + passwd
   154  	}
   155  	return fmt.Sprintf("postgres://%s%s@%s", e.Master.Username, passwd, e.Slave.Host())
   156  }
   157  
   158  func (e Endpoint) Name() string { return "postgres-cli" }
   159  
   160  func SwitchSlave(db sqlx.DBExecutor) sqlx.DBExecutor {
   161  	if slave, ok := db.(CanSlave); ok {
   162  		return slave.UseSlave()
   163  	}
   164  	return db
   165  }
   166  
   167  type CanSlave interface {
   168  	UseSlave() sqlx.DBExecutor
   169  }