generated from amosproj/amos202Xss0Y-projname
/
setup.go
105 lines (84 loc) · 2.72 KB
/
setup.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
package setup
import (
"database/sql"
"fmt"
"os"
"strings"
"time"
"github.com/uptrace/bun"
"github.com/uptrace/bun/dialect/pgdialect"
"github.com/uptrace/bun/driver/pgdriver"
"github.com/uptrace/bun/extra/bundebug"
"k8s.io/klog"
)
const (
connectionTimeout = 5 * time.Second
dialTimeout = 5 * time.Second
retryInterval = 2 * time.Second // retry interval for testDBConnection
retryTimeout = 30 * time.Second // final timeout for testDBConnection
)
// DBConnection setup database connection.
func DBConnection() *bun.DB {
dbUser, exists := os.LookupEnv("DB_USER")
if !exists {
klog.Warning("DB_USER environment variable is not set. Trying dbUser = postgres")
dbUser = "postgres"
}
dbPassword, exists := os.LookupEnv("DB_PASSWORD")
if !exists {
klog.Warning("DB_PASSWORD environment variable is not set. Trying dbPassword = example")
dbPassword = "example"
}
dbURI, exists := os.LookupEnv("DB_URI")
if !exists {
klog.Warning("DB_URI environment variable is not set. Trying dbURI = localhost")
dbURI = "localhost"
}
pgAddr := fmt.Sprintf("%s:5432", dbURI)
pgconn := pgdriver.NewConnector(
pgdriver.WithNetwork("tcp"),
pgdriver.WithAddr(pgAddr),
pgdriver.WithUser(dbUser),
pgdriver.WithPassword(dbPassword),
pgdriver.WithDatabase("postgres"),
pgdriver.WithInsecure(true),
pgdriver.WithTimeout(connectionTimeout),
pgdriver.WithDialTimeout(dialTimeout),
)
sqldb := sql.OpenDB(pgconn)
dbConn := bun.NewDB(sqldb, pgdialect.New())
// This line can toggle verbose debugging output for every query that is send
// db.AddQueryHook(bundebug.NewQueryHook(bundebug.WithVerbose(true)))
dbConn.AddQueryHook(bundebug.NewQueryHook())
testDBConnection(dbConn)
return dbConn
}
// TestDBConnection attempts to establish a connection with the database and retries for 30 seconds if unsuccessful.
func testDBConnection(db *bun.DB) {
start := time.Now()
for {
err := db.Ping()
if err != nil {
if time.Since(start) < retryTimeout {
// Check if the error is due to database not started
switch {
case strings.Contains(err.Error(), "connection refused"):
// Check if the error is due to database not started
klog.Warningf("Database is not up yet: %v. Retrying...", err)
case strings.Contains(err.Error(), "authentication failed"):
// Check if the error is due to wrong authentication data
klog.Fatalf("Wrong database authentication data: %v", err)
default:
// Other database error
klog.Warningf("Could not connect to database: %v. Retrying...", err)
}
time.Sleep(retryInterval)
} else {
klog.Fatalf("Could not connect to database after 30 seconds: %v", err)
}
} else {
klog.Info("Successfully connected to database.")
return
}
}
}