101 lines
1.8 KiB
Go
101 lines
1.8 KiB
Go
package headscale
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
const dbVersion = "1"
|
|
|
|
// KV is a key-value store in a psql table. For future use...
|
|
type KV struct {
|
|
Key string
|
|
Value string
|
|
}
|
|
|
|
func (h *Headscale) initDB() error {
|
|
db, err := h.db()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if h.dbType == "postgres" {
|
|
db.Exec("create extension if not exists \"uuid-ossp\";")
|
|
}
|
|
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
|
|
}
|
|
|
|
err = h.setValue("db_version", dbVersion)
|
|
return err
|
|
}
|
|
|
|
func (h *Headscale) db() (*gorm.DB, error) {
|
|
var db *gorm.DB
|
|
var err error
|
|
switch h.dbType {
|
|
case "sqlite3":
|
|
db, err = gorm.Open(sqlite.Open(h.dbString), &gorm.Config{
|
|
DisableForeignKeyConstraintWhenMigrating: true,
|
|
})
|
|
case "postgres":
|
|
db, err = gorm.Open(postgres.Open(h.dbString), &gorm.Config{
|
|
DisableForeignKeyConstraintWhenMigrating: true,
|
|
})
|
|
}
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if h.dbDebug {
|
|
db.Debug()
|
|
}
|
|
return db, nil
|
|
}
|
|
|
|
func (h *Headscale) getValue(key string) (string, error) {
|
|
db, err := h.db()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
var row KV
|
|
if result := db.First(&row, "key = ?", key); errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
|
return "", errors.New("not found")
|
|
}
|
|
return row.Value, nil
|
|
}
|
|
|
|
func (h *Headscale) setValue(key string, value string) error {
|
|
kv := KV{
|
|
Key: key,
|
|
Value: value,
|
|
}
|
|
db, err := h.db()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = h.getValue(key)
|
|
if err == nil {
|
|
db.Model(&kv).Where("key = ?", key).Update("value", value)
|
|
return nil
|
|
}
|
|
|
|
db.Create(kv)
|
|
return nil
|
|
}
|