/
connect.go
86 lines (72 loc) · 2.44 KB
/
connect.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
package postgres
import (
"github.com/jackc/pgx"
"github.com/pkg/errors"
"github.com/wal-g/tracelog"
)
// Connect establishes a connection to postgres using
// a UNIX socket. Must export PGHOST and run with `sudo -E -u postgres`.
// If PGHOST is not set or if the connection fails, an error is returned
// and the connection is `<nil>`.
//
// Example: PGHOST=/var/run/postgresql or PGHOST=10.0.0.1
func Connect(configOptions ...func(config *pgx.ConnConfig) error) (*pgx.Conn, error) {
config, err := pgx.ParseEnvLibpq()
if err != nil {
return nil, errors.Wrap(err, "Connect: unable to read environment variables")
}
// apply passed custom config options, if any
for _, option := range configOptions {
err := option(&config)
if err != nil {
return nil, err
}
}
conn, err := pgx.Connect(config)
if err != nil {
conn, err = tryConnectToGpSegment(config)
if err != nil && config.Host != "localhost" {
tracelog.ErrorLogger.Println(err.Error())
tracelog.ErrorLogger.Println("Failed to connect using provided PGHOST and PGPORT, trying localhost:5432")
config.Host = "localhost"
config.Port = 5432
conn, err = pgx.Connect(config)
}
if err != nil {
return nil, errors.Wrap(err, "Connect: postgres connection failed")
}
}
var archiveMode string
// TODO: Move this logic to queryRunner
err = conn.QueryRow("show archive_mode").Scan(&archiveMode)
if err != nil {
return nil, errors.Wrap(err, "Connect: postgres archive_mode test failed")
}
if archiveMode != "on" && archiveMode != "always" {
tracelog.WarningLogger.Println(
"It seems your archive_mode is not enabled. This will cause inconsistent backup. " +
"Please consider configuring WAL archiving.")
} else {
var archiveCommand string
err = conn.QueryRow("show archive_command").Scan(&archiveCommand)
if err != nil {
return nil, errors.Wrap(err, "Connect: postgres archive_mode test failed")
}
if len(archiveCommand) == 0 || archiveCommand == "(disabled)" {
tracelog.WarningLogger.Println(
"It seems your archive_command is not configured. This will cause inconsistent backup." +
" Please consider configuring WAL archiving.")
}
}
return conn, nil
}
// nolint:gocritic
func tryConnectToGpSegment(config pgx.ConnConfig) (*pgx.Conn, error) {
config.RuntimeParams["gp_role"] = "utility"
conn, err := pgx.Connect(config)
if err != nil {
config.RuntimeParams["gp_session_role"] = "utility"
conn, err = pgx.Connect(config)
}
return conn, err
}