/
postgresql.go
121 lines (104 loc) · 2.73 KB
/
postgresql.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
package probes
import (
"context"
"crypto/tls"
"crypto/x509"
"errors"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
var _ Prober = PostgreSQL{}
// PostgreSQL is a prober to test the uptime of a PostgreSQL instance.
type PostgreSQL struct {
db *pgxpool.Pool
Service string
Name string
Namespace string
Organization string
HighAvailable bool
ServiceLevel string
}
// Close closes open connections to the PostgreSQL server.
func (p PostgreSQL) Close() error {
if p.db != nil {
p.db.Close()
}
return nil
}
// GetInfo returns the prober infos
func (p PostgreSQL) GetInfo() ProbeInfo {
return ProbeInfo{
Service: p.Service,
Name: p.Name,
Namespace: p.Namespace,
Organization: p.Organization,
HighAvailable: p.HighAvailable,
ServiceLevel: p.ServiceLevel,
}
}
// Probe sends a test query to the configured PostgreSQL server.
// Will return an error if the prober does not have a valid db connection.
func (p PostgreSQL) Probe(ctx context.Context) error {
if p.db == nil {
return errors.New("invalid credentials")
}
_, err := p.db.Exec(ctx, "SELECT 1")
return err
}
// NewPostgreSQL connects to the provided dsn and returns a prober
func NewPostgreSQL(service, name, namespace, dsn, organization, sla string, ha bool, ops ...func(*pgxpool.Config) error) (*PostgreSQL, error) {
conf, err := pgxpool.ParseConfig(dsn)
if err != nil {
return nil, err
}
conf.ConnConfig.ConnectTimeout = 5 * time.Second
conf.MaxConns = 1
for _, op := range ops {
err := op(conf)
if err != nil {
return nil, err
}
}
db, err := pgxpool.NewWithConfig(context.Background(), conf)
if err != nil {
return nil, err
}
return &PostgreSQL{
db: db,
Service: service,
Name: name,
Namespace: namespace,
Organization: organization,
HighAvailable: ha,
ServiceLevel: sla,
}, nil
}
// NewFailingPostgreSQL creates a prober that will fail.
// Can be used if the controller can't access valid credentials.
func NewFailingPostgreSQL(service, name, namespace string) (*PostgreSQL, error) {
return &PostgreSQL{
Service: service,
Name: name,
Namespace: namespace,
}, nil
}
// PGWithCA adds the provided CA to the rootCAs of the pgxpool.
func PGWithCA(ca []byte) func(*pgxpool.Config) error {
return func(conf *pgxpool.Config) error {
if conf.ConnConfig.TLSConfig == nil {
conf.ConnConfig.TLSConfig = &tls.Config{
RootCAs: x509.NewCertPool(),
}
}
if conf.ConnConfig.TLSConfig.RootCAs == nil {
conf.ConnConfig.TLSConfig.RootCAs = x509.NewCertPool()
}
if ca == nil {
return errors.New("got nil CA")
}
if !conf.ConnConfig.TLSConfig.RootCAs.AppendCertsFromPEM(ca) {
return errors.New("cannot append root CA certificates")
}
return nil
}
}