github.com/profzone/eden-framework@v1.0.10/pkg/client/postgre/postgres.go (about)

     1  package confpostgres
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"github.com/profzone/eden-framework/pkg/sqlx"
     7  	"github.com/profzone/eden-framework/pkg/sqlx/postgresqlconnector"
     8  	"github.com/profzone/envconfig"
     9  	"time"
    10  )
    11  
    12  type Postgres struct {
    13  	Host            string
    14  	SlaveHost       string
    15  	Port            int
    16  	User            string
    17  	Password        envconfig.Password
    18  	Extra           string
    19  	Extensions      []string
    20  	PoolSize        int
    21  	ConnMaxLifetime envconfig.Duration
    22  	Database        *sqlx.Database `ignored:"true"`
    23  
    24  	*sqlx.DB `ignored:"true"`
    25  	slaveDB  *sqlx.DB `ignored:"true"`
    26  }
    27  
    28  func (m *Postgres) LivenessCheck() map[string]string {
    29  	s := map[string]string{}
    30  
    31  	_, err := m.DB.ExecContext(context.Background(), "SELECT 1")
    32  	if err != nil {
    33  		s[m.Host] = err.Error()
    34  	} else {
    35  		s[m.Host] = "ok"
    36  	}
    37  
    38  	if m.slaveDB != nil {
    39  		_, err := m.slaveDB.ExecContext(context.Background(), "SELECT 1")
    40  		if err != nil {
    41  			s[m.SlaveHost] = err.Error()
    42  		} else {
    43  			s[m.SlaveHost] = "ok"
    44  		}
    45  	}
    46  
    47  	return s
    48  }
    49  
    50  func (m *Postgres) SetDefaults() {
    51  	if m.Host == "" {
    52  		m.Host = "127.0.0.1"
    53  	}
    54  
    55  	if m.Port == 0 {
    56  		m.Port = 5432
    57  	}
    58  
    59  	if m.PoolSize == 0 {
    60  		m.PoolSize = 10
    61  	}
    62  
    63  	if m.ConnMaxLifetime == 0 {
    64  		m.ConnMaxLifetime = envconfig.Duration(1 * time.Hour)
    65  	}
    66  
    67  	if m.Extra == "" {
    68  		m.Extra = "sslmode=disable"
    69  	}
    70  }
    71  
    72  func (m *Postgres) url(host string) string {
    73  	password := m.Password
    74  	if password != "" {
    75  		password = ":" + password
    76  	}
    77  	return fmt.Sprintf("postgres://%s%s@%s:%d", m.User, password, host, m.Port)
    78  }
    79  
    80  func (m *Postgres) conn(host string) (*sqlx.DB, error) {
    81  	db := m.Database.OpenDB(&postgresqlconnector.PostgreSQLConnector{
    82  		Host:       m.url(host),
    83  		Extra:      m.Extra,
    84  		Extensions: m.Extensions,
    85  	})
    86  
    87  	db.SetMaxOpenConns(m.PoolSize)
    88  	db.SetMaxIdleConns(m.PoolSize / 2)
    89  	db.SetConnMaxLifetime(time.Duration(m.ConnMaxLifetime))
    90  
    91  	_, err := db.ExecContext(context.Background(), "SELECT 1")
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  
    96  	return db, nil
    97  }
    98  
    99  func (m *Postgres) UseSlave() sqlx.DBExecutor {
   100  	if m.slaveDB != nil {
   101  		return m.slaveDB
   102  	}
   103  	return m.DB
   104  }
   105  
   106  func (m *Postgres) Init() {
   107  	r := Retry{Repeats: 5, Interval: envconfig.Duration(1 * time.Second)}
   108  
   109  	err := r.Do(func() error {
   110  		db, err := m.conn(m.Host)
   111  		if err != nil {
   112  			return err
   113  		}
   114  		m.DB = db
   115  		return nil
   116  	})
   117  
   118  	if err != nil {
   119  		panic(err)
   120  	}
   121  
   122  	if m.SlaveHost != "" {
   123  		err := r.Do(func() error {
   124  			db, err := m.conn(m.Host)
   125  			if err != nil {
   126  				return err
   127  			}
   128  			m.slaveDB = db
   129  			return nil
   130  		})
   131  
   132  		if err != nil {
   133  			panic(err)
   134  		}
   135  	}
   136  }
   137  
   138  func SwitchSlave(executor sqlx.DBExecutor) sqlx.DBExecutor {
   139  	if canSlave, ok := executor.(CanSlave); !ok {
   140  		return canSlave.UseSlave()
   141  	}
   142  	return executor
   143  }
   144  
   145  type CanSlave interface {
   146  	UseSlave() sqlx.DBExecutor
   147  }