Passed
Push — master ( 945ad4...699fd0 )
by Tolga
01:30 queued 36s
created

postgres.*Postgres.IsReady   A

Complexity

Conditions 2

Size

Total Lines 7
Code Lines 6

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 2
eloc 6
nop 1
dl 0
loc 7
rs 10
c 0
b 0
f 0
1
package postgres
2
3
import (
4
	"context"
5
	"time"
6
7
	"github.com/jackc/pgx/v5"
8
	"github.com/jackc/pgx/v5/pgxpool"
9
10
	"github.com/Masterminds/squirrel"
11
)
12
13
// Postgres - Structure for Postresql instance
14
type Postgres struct {
15
	ReadPool  *pgxpool.Pool
16
	WritePool *pgxpool.Pool
17
18
	Builder squirrel.StatementBuilderType
19
	// options
20
	maxDataPerWrite       int
21
	maxRetries            int
22
	watchBufferSize       int
23
	maxConnectionLifeTime time.Duration
24
	maxConnectionIdleTime time.Duration
25
	maxOpenConnections    int
26
	maxIdleConnections    int
27
	simpleMode            bool
28
}
29
30
// New - Creates new postgresql db instance
31
func New(uri string, opts ...Option) (*Postgres, error) {
32
	pg := &Postgres{
33
		maxOpenConnections: _defaultMaxOpenConnections,
34
		maxIdleConnections: _defaultMaxIdleConnections,
35
		maxDataPerWrite:    _defaultMaxDataPerWrite,
36
		maxRetries:         _defaultMaxRetries,
37
		watchBufferSize:    _defaultWatchBufferSize,
38
		simpleMode:         _defaultSimpleMode,
39
	}
40
41
	// Custom options
42
	for _, opt := range opts {
43
		opt(pg)
44
	}
45
46
	pg.Builder = squirrel.StatementBuilder.PlaceholderFormat(squirrel.Dollar)
47
48
	writeConfig, err := pgxpool.ParseConfig(uri)
49
	if err != nil {
50
		return nil, err
51
	}
52
53
	readConfig, err := pgxpool.ParseConfig(uri)
54
	if err != nil {
55
		return nil, err
56
	}
57
58
	if pg.simpleMode {
59
		writeConfig.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol
60
		readConfig.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol
61
	}
62
63
	writeConfig.MinConns = int32(pg.maxIdleConnections)
64
	readConfig.MinConns = int32(pg.maxIdleConnections)
65
66
	writeConfig.MaxConns = int32(pg.maxOpenConnections)
67
	readConfig.MaxConns = int32(pg.maxOpenConnections)
68
69
	writeConfig.MaxConnIdleTime = pg.maxConnectionIdleTime
70
	readConfig.MaxConnIdleTime = pg.maxConnectionIdleTime
71
72
	writeConfig.MaxConnLifetime = pg.maxConnectionLifeTime
73
	readConfig.MaxConnLifetime = pg.maxConnectionLifeTime
74
75
	writeConfig.MaxConnLifetimeJitter = time.Duration(0.2 * float64(pg.maxConnectionLifeTime))
76
	readConfig.MaxConnLifetimeJitter = time.Duration(0.2 * float64(pg.maxConnectionLifeTime))
77
78
	if _, ok := readConfig.ConnConfig.Config.RuntimeParams["plan_cache_mode"]; !ok {
79
		readConfig.ConnConfig.Config.RuntimeParams["plan_cache_mode"] = "force_custom_plan"
80
	}
81
82
	if _, ok := writeConfig.ConnConfig.Config.RuntimeParams["plan_cache_mode"]; !ok {
83
		writeConfig.ConnConfig.Config.RuntimeParams["plan_cache_mode"] = "force_custom_plan"
84
	}
85
86
	initialContext, cancelInit := context.WithTimeout(context.Background(), 5*time.Second)
87
	defer cancelInit()
88
89
	pg.WritePool, err = pgxpool.NewWithConfig(initialContext, writeConfig)
90
	if err != nil {
91
		return nil, err
92
	}
93
	pg.ReadPool, err = pgxpool.NewWithConfig(initialContext, readConfig)
94
	if err != nil {
95
		return nil, err
96
	}
97
98
	return pg, nil
99
}
100
101
func (p *Postgres) GetMaxDataPerWrite() int {
102
	return p.maxDataPerWrite
103
}
104
105
func (p *Postgres) GetMaxRetries() int {
106
	return p.maxRetries
107
}
108
109
func (p *Postgres) GetWatchBufferSize() int {
110
	return p.watchBufferSize
111
}
112
113
// GetEngineType - Get the engine type which is postgresql in string
114
func (p *Postgres) GetEngineType() string {
115
	return "postgres"
116
}
117
118
// Close - Close postgresql instance
119
func (p *Postgres) Close() error {
120
	p.ReadPool.Close()
121
	p.WritePool.Close()
122
	return nil
123
}
124
125
// IsReady - Check if database is ready
126
func (p *Postgres) IsReady(ctx context.Context) (bool, error) {
127
	ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
128
	defer cancel()
129
	if err := p.ReadPool.Ping(ctx); err != nil {
130
		return false, err
131
	}
132
	return true, nil
133
}
134