/
db.go
97 lines (85 loc) · 2.08 KB
/
db.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
package database
import (
"fmt"
"io"
"strings"
"time"
"github.com/charmbracelet/log"
"github.com/flamego/flamego"
"github.com/pkg/errors"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/logger"
"github.com/pgrok/pgrok/internal/conf"
)
// DB is the database handle.
type DB struct {
*gorm.DB
}
// New returns a new database handle with given configuration.
func New(logWriter io.Writer, config *conf.Database) (*DB, error) {
if config == nil {
return nil, errors.New("no database config provided")
}
level := logger.Info
if flamego.Env() == flamego.EnvTypeProd {
level = logger.Warn
}
// NOTE: AutoMigrate does not respect logger passed in gorm.Config.
logger.Default = logger.New(
&gormLogger{
Logger: log.NewWithOptions(
logWriter,
log.Options{
TimeFormat: time.DateTime,
Level: log.DebugLevel,
Prefix: "gorm",
ReportTimestamp: true,
},
),
},
logger.Config{
SlowThreshold: 1000 * time.Millisecond,
LogLevel: level,
},
)
dsn := fmt.Sprintf(
"user='%s' password='%s' host='%s' port='%d' dbname='%s' search_path='public' application_name='pgrokd'",
config.User, config.Password, config.Host, config.Port, config.Database,
)
db, err := gorm.Open(
postgres.Open(dsn),
&gorm.Config{
SkipDefaultTransaction: true,
NowFunc: func() time.Time {
return time.Now().UTC().Truncate(time.Microsecond)
},
},
)
if err != nil {
return nil, errors.Wrap(err, "open database")
}
sqlDB, err := db.DB()
if err != nil {
return nil, errors.Wrap(err, "get underlying *sql.DB")
}
sqlDB.SetMaxOpenConns(30)
sqlDB.SetMaxIdleConns(30)
sqlDB.SetConnMaxLifetime(time.Minute)
err = db.AutoMigrate(&Principal{}, &HostKey{})
if err != nil {
return nil, errors.Wrap(err, "auto migrate")
}
return &DB{db}, nil
}
// gormLogger is a wrapper of io.Writer for the GORM's logger.Writer.
type gormLogger struct {
*log.Logger
}
func (l *gormLogger) Printf(format string, args ...any) {
print := l.Debug
if strings.Contains(format, "[error]") {
print = l.Error
}
print(fmt.Sprintf(format, args...))
}