/
database.go
93 lines (77 loc) · 2.42 KB
/
database.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
package database
import (
"context"
"database/sql"
"fmt"
"time"
"entgo.io/ent/dialect"
entsql "entgo.io/ent/dialect/sql"
"github.com/database64128/proxy-sharing-go/ent"
"github.com/database64128/proxy-sharing-go/ent/migrate"
"github.com/database64128/proxy-sharing-go/jsonhelper"
_ "github.com/go-sql-driver/mysql"
_ "github.com/jackc/pgx/v5/stdlib"
"go.uber.org/zap"
)
// Config is the configuration for the database.
type Config struct {
// Driver is the database driver.
Driver string `json:"driver"`
// DSN is the data source name.
DSN string `json:"dsn"`
// MaxOpenConns is the maximum number of open connections to the database.
MaxOpenConns int `json:"maxOpenConns"`
// MaxIdleConns is the maximum number of connections in the idle connection pool.
MaxIdleConns int `json:"maxIdleConns"`
// ConnMaxLifetime is the maximum amount of time a connection may be reused.
ConnMaxLifetime jsonhelper.Duration `json:"connMaxLifetime"`
// ConnMaxIdleTime is the maximum amount of time a connection may be idle.
ConnMaxIdleTime jsonhelper.Duration `json:"connMaxIdleTime"`
// Debug enables verbose logging.
Debug bool `json:"debug"`
// NoAutoMigrate disables auto-migration.
NoAutoMigrate bool `json:"noAutoMigrate"`
}
// Open opens the database and runs auto-migration.
func (c *Config) Open(ctx context.Context, logger *zap.Logger) (*ent.Client, error) {
var (
db *sql.DB
err error
)
switch c.Driver {
case dialect.MySQL:
db, err = sql.Open(c.Driver, c.DSN)
case dialect.SQLite:
db, err = openSQLiteDB(c.DSN)
case dialect.Postgres:
db, err = sql.Open("pgx", c.DSN)
default:
return nil, fmt.Errorf("unsupported driver: %q", c.Driver)
}
if err != nil {
return nil, fmt.Errorf("failed to open database: %w", err)
}
db.SetMaxOpenConns(c.MaxOpenConns)
db.SetMaxIdleConns(c.MaxIdleConns)
db.SetConnMaxLifetime(time.Duration(c.ConnMaxLifetime))
db.SetConnMaxIdleTime(time.Duration(c.ConnMaxIdleTime))
drv := entsql.OpenDB(c.Driver, db)
var opts []ent.Option
if !c.Debug {
opts = []ent.Option{ent.Driver(drv)}
} else {
sugar := logger.Sugar()
opts = []ent.Option{ent.Driver(drv), ent.Debug(), ent.Log(sugar.Debugln)}
}
client := ent.NewClient(opts...)
if !c.NoAutoMigrate {
if err = client.Schema.Create(
ctx,
migrate.WithDropColumn(true),
migrate.WithDropIndex(true),
); err != nil {
return nil, fmt.Errorf("failed to create schema resources: %w", err)
}
}
return client, nil
}