fyp/logic/utils/config.go

package utils

import (
	"database/sql"
	"os"
	"strings"

	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/tasks/utils"

	toml "github.com/BurntSushi/toml"
	"github.com/charmbracelet/log"
)

type WorkerConfig struct {
	NumberOfWorkers int    `toml:"number_of_workers"`
	Pulling         string `toml:"pulling_time"`
}

type ServiceUser struct {
	User   string `toml:"user"`
	UserId string `toml:"__user__id__"`
}

type DbInfo struct {
	MaxConnections int `toml:"max_connections"`
}

type Config struct {
	Hostname        string
	Port            int
	NumberOfWorkers int `toml:"number_of_workers"`

	SupressCuda      int `toml:"supress_cuda"`
	CleanUpOnStartup int `toml:"clean_up_on_startup"`

	GpuWorker WorkerConfig `toml:"Worker"`

	ServiceUser ServiceUser `toml:"ServiceUser"`

	DbInfo DbInfo `toml:"DB"`
}

func LoadConfig() Config {

	log.Info("Loading the config file")

	dat, err := os.ReadFile("./config.toml")
	if err != nil {
		log.Error("Failed to load config file", "err", err)
		// Use default values
		return Config{
			Hostname:         "localhost",
			Port:             8000,
			NumberOfWorkers:  10,
			CleanUpOnStartup: 1,
			SupressCuda:      1,
			GpuWorker: WorkerConfig{
				NumberOfWorkers: 1,
				Pulling:         "500ms",
			},
			ServiceUser: ServiceUser{
				User:   "Service",
				UserId: "",
			},
			DbInfo: DbInfo{
				MaxConnections: 200,
			},
		}
	}

	var conf Config
	_, err = toml.Decode(string(dat), &conf)

	if conf.SupressCuda == 1 {
		log.Warn("Supressing Cuda Messages!")
		os.Setenv("TF_CPP_MIN_VLOG_LEVEL", "3")
		os.Setenv("TF_CPP_MIN_LOG_LEVEL", "3")
	}

	return conf
}

func failLog(err error) {
	if err == nil {
		return
	}
	log.Fatal("Failed on setup", "error", err)
}

func (c *Config) Cleanup(db *sql.DB) {
	if c.CleanUpOnStartup != 1 {
		return
	}

	_, err := db.Exec("update models set status=$1 where status=$2", FAILED_PREPARING_ZIP_FILE, PREPARING_ZIP_FILE)
	failLog(err)
	_, err = db.Exec("update models set status=$1 where status=$2", FAILED_PREPARING, PREPARING)
	failLog(err)
	_, err = db.Exec("update tasks set status=$1 where status=$2", TASK_PICKED_UP, TASK_TODO)
	failLog(err)

	tasks, err := GetDbMultitple[Task](db, "tasks where status=$1", TASK_RUNNING)
	failLog(err)

	base := BasePackStruct{Db: db, Logger: log.Default()}

	for i := range tasks {
		if tasks[i].TaskType == int(TASK_TYPE_CLASSIFICATION) {
			tasks[i].UpdateStatus(base, TASK_TODO, "Reseting task")
			continue
		}
		if tasks[i].TaskType == int(TASK_TYPE_RETRAINING) {
			tasks[i].UpdateStatus(base, TASK_FAILED_RUNNING, "Task inturupted by server restart please try again")
			_, err = db.Exec("update models set status=$1 where id=$2", READY_RETRAIN_FAILED, tasks[i].ModelId)
			failLog(err)
			continue
		}
		if tasks[i].TaskType == int(TASK_TYPE_TRAINING) {
			tasks[i].UpdateStatus(base, TASK_FAILED_RUNNING, "Task inturupted by server restart please try again")
			_, err = db.Exec("update models set status=$1 where id=$2", FAILED_TRAINING, tasks[i].ModelId)
			failLog(err)
			continue
		}
	}
}

func (c *Config) GenerateToken(db *sql.DB) {
	if c.ServiceUser.User == "" {
		log.Fatal("A user needs to be set in a configuration file")
	}

	var user struct {
		Password string `db:"password"`
		UserId   string `db:"id"`
	}

	err := GetDBOnce(db, &user, "users where username=$1;", c.ServiceUser.User)
	if err == NotFoundError {
		var newUser = struct {
			Username string
			Email    string
			Salt     string
			Password string
			UserType UserType `db:"user_type"`
		}{
			c.ServiceUser.User,
			c.ServiceUser.User,
			"",
			"",
			User_Admin,
		}
		id, err := InsertReturnId(db, &newUser, "users", "id")
		if err != nil {
			log.Fatal("Failed to create user", "err", err)
		}
		c.ServiceUser.UserId = id
	} else if err != nil {
		log.Fatal("To get user name", "err", err)
		return
	} else {
		if len(strings.ReplaceAll(user.Password, " ", "")) != 0 {
			log.Fatal("User already exists and is not the service user", "user", user)
		}
		c.ServiceUser.UserId = user.UserId
	}
}