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 := strings.TrimPrefix(cfg.DatabaseURL, "sqlite://")
	if dsn == cfg.DatabaseURL {
		dsn = cfg.DatabaseURL
	}

	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{},
	)
}