headscale/db.go

96 lines
1.7 KiB
Go
Raw Normal View History

2020-06-21 04:32:08 -06:00
package headscale
import (
"errors"
2021-06-24 07:44:19 -06:00
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
2020-06-21 04:32:08 -06:00
)
const dbVersion = "1"
2021-02-23 12:11:30 -07:00
// KV is a key-value store in a psql table. For future use...
2020-06-21 04:32:08 -06:00
type KV struct {
Key string
Value string
}
func (h *Headscale) initDB() error {
2021-07-04 13:40:46 -06:00
db, err := h.openDB()
2020-06-21 04:32:08 -06:00
if err != nil {
return err
}
2021-07-04 13:40:46 -06:00
h.db = db
if h.dbType == "postgres" {
db.Exec("create extension if not exists \"uuid-ossp\";")
}
2021-06-24 07:44:19 -06:00
err = db.AutoMigrate(&Machine{})
if err != nil {
return err
}
err = db.AutoMigrate(&KV{})
if err != nil {
return err
}
err = db.AutoMigrate(&Namespace{})
if err != nil {
return err
}
err = db.AutoMigrate(&PreAuthKey{})
if err != nil {
return err
}
2020-06-21 04:32:08 -06:00
err = h.setValue("db_version", dbVersion)
return err
2020-06-21 04:32:08 -06:00
}
2021-07-04 13:40:46 -06:00
func (h *Headscale) openDB() (*gorm.DB, error) {
2021-06-24 07:44:19 -06:00
var db *gorm.DB
var err error
switch h.dbType {
case "sqlite3":
db, err = gorm.Open(sqlite.Open(h.dbString), &gorm.Config{
DisableForeignKeyConstraintWhenMigrating: true,
})
2021-06-24 07:44:19 -06:00
case "postgres":
db, err = gorm.Open(postgres.Open(h.dbString), &gorm.Config{
DisableForeignKeyConstraintWhenMigrating: true,
})
2021-06-24 07:44:19 -06:00
}
2020-06-21 04:32:08 -06:00
if err != nil {
return nil, err
}
if h.dbDebug {
2021-06-24 07:44:19 -06:00
db.Debug()
}
2020-06-21 04:32:08 -06:00
return db, nil
}
func (h *Headscale) getValue(key string) (string, error) {
var row KV
2021-07-04 13:40:46 -06:00
if result := h.db.First(&row, "key = ?", key); errors.Is(result.Error, gorm.ErrRecordNotFound) {
2020-06-21 04:32:08 -06:00
return "", errors.New("not found")
}
return row.Value, nil
}
func (h *Headscale) setValue(key string, value string) error {
kv := KV{
Key: key,
Value: value,
}
2021-07-04 13:40:46 -06:00
_, err := h.getValue(key)
2020-06-21 04:32:08 -06:00
if err == nil {
2021-07-04 13:40:46 -06:00
h.db.Model(&kv).Where("key = ?", key).Update("value", value)
2020-06-21 04:32:08 -06:00
return nil
}
2021-07-04 13:40:46 -06:00
h.db.Create(kv)
2020-06-21 04:32:08 -06:00
return nil
}