-
Notifications
You must be signed in to change notification settings - Fork 4
/
database.go
114 lines (101 loc) · 2.92 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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
package cmd
import (
"errors"
"flag"
"fmt"
"regexp"
"github.com/efixler/envflags"
"github.com/efixler/scrape/internal/storage/mysql"
"github.com/efixler/scrape/internal/storage/sqlite"
"github.com/efixler/scrape/store"
)
var (
DefaultDatabase = DatabaseSpec{Type: "sqlite", Path: "scrape_data/scrape.db"}
ErrDatabaseFormat = errors.New("database spec must be in the format <type>:<path spec>")
ErrUnsupportedDatabase = errors.New("unsupported database type")
dsnRegex = regexp.MustCompile(`^(\w{5,10}):(.+)`)
)
type DatabaseSpec struct {
Type string
Path string
}
func (d DatabaseSpec) String() string {
return d.Type + ":" + d.Path
}
// Format: <type>:<path spec>
// Example: sqlite:scrape_data/scrape.db
// Example: sqlite::memory:
func NewDatabaseSpec(s string) (DatabaseSpec, error) {
matches := dsnRegex.FindStringSubmatch(s)
if matches == nil {
return DatabaseSpec{}, ErrDatabaseFormat
}
spec := DatabaseSpec{
Type: matches[1],
Path: matches[2],
}
return spec, nil
}
func NewDatabaseValue(env string, def DatabaseSpec) *envflags.Value[DatabaseSpec] {
converter := NewDatabaseSpec
val := envflags.NewEnvFlagValue(env, def, converter)
return val
}
type DatabaseFlags struct {
database *envflags.Value[DatabaseSpec]
username *envflags.Value[string]
password *envflags.Value[string]
Migrate bool
}
func AddDatabaseFlags(baseEnv string, flags *flag.FlagSet, migrateFlag bool) *DatabaseFlags {
dbFlags := &DatabaseFlags{
database: NewDatabaseValue(baseEnv, DefaultDatabase),
username: envflags.NewString(baseEnv+"_USER", ""),
password: envflags.NewString(baseEnv+"_PASSWORD", ""),
}
dbFlags.database.AddTo(flags, "database", "Database type:path")
dbFlags.username.AddTo(flags, "db-user", "Database user")
dbFlags.password.AddTo(flags, "db-password", "Database password")
if migrateFlag {
flags.BoolVar(
&dbFlags.Migrate,
"migrate",
false,
"Migrate the database to the latest version (creating if necessary)",
)
}
return dbFlags
}
func (f DatabaseFlags) String() DatabaseSpec {
return f.database.Get()
}
func (f DatabaseFlags) Database() (store.Factory, error) {
return database(f.database.Get(), f.username.Get(), f.password.Get(), f.Migrate)
}
func (f DatabaseFlags) MustDatabase() store.Factory {
dbF, err := f.Database()
if err != nil {
panic(fmt.Sprintf("error making database factory from flags: %v", err))
}
return dbF
}
func database(spec DatabaseSpec, username string, password string, noSchema bool) (store.Factory, error) {
switch spec.Type {
case "sqlite3":
fallthrough
case "sqlite":
return sqlite.Factory(sqlite.File(spec.Path)), nil
case "mysql":
options := []mysql.Option{
mysql.NetAddress(spec.Path),
mysql.Username(username),
mysql.Password(password),
}
if noSchema {
options = append(options, mysql.ForMigration())
}
return mysql.Factory(options...), nil
default:
return nil, errors.New("no implementation for " + spec.Type)
}
}