llama-link/internal/db/db.go

package db

import (
	"log/slog"
	"strings"
	"time"

	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/llamalink/llamalink/internal/config"
)

func Open(cfg *config.Config) (*gorm.DB, error) {
	dsn := cfg.DatabaseURL
	if strings.HasPrefix(dsn, "sqlite://") {
		dsn = strings.TrimPrefix(dsn, "sqlite://")
		if strings.HasPrefix(dsn, "/") {
			dsn = dsn[1:]
		}
	}

	gormConfig := &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	}

	db, err := gorm.Open(sqlite.Open(dsn), gormConfig)
	if err != nil {
		return nil, err
	}

	sqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}

	sqlDB.SetMaxOpenConns(cfg.DatabaseMaxOpenConns)
	sqlDB.SetMaxIdleConns(cfg.DatabaseMaxIdleConns)
	sqlDB.SetConnMaxLifetime(time.Duration(cfg.DatabaseConnMaxLifetime) * time.Second)

	return db, nil
}

func Migrate(db *gorm.DB) error {
	slog.Info("running database migrations")
	return db.AutoMigrate(
		&ApiKey{},
		&Model{},
		&UsageLog{},
		&Quota{},
		&Webhook{},
	)
}