database.go

package main

import (
	"time"

	"github.com/glebarez/sqlite"
	"gorm.io/gorm"
)

type DB struct {
	*gorm.DB
}

type Secret struct {
	ID            uint           `gorm:"primaryKey;autoIncrement" json:"id"`
	Slug          string         `gorm:"type:varchar(32);uniqueIndex:IX_secrets_slug;not null" json:"slug"`
	Secret        *string        `gorm:"type:text" json:"secret,omitempty"`
	FileName      *string        `gorm:"type:varchar(256)" json:"file_name,omitempty"`
	FileSize      *int64         `gorm:"type:integer" json:"file_size,omitempty"`
	FileExtension *string        `gorm:"type:varchar(10)" json:"file_extension,omitempty"`
	FileMime      *string        `gorm:"type:varchar(100)" json:"file_mime,omitempty"`
	FileContents  *string        `gorm:"type:text" json:"file_contents,omitempty"`
	CreatedAt     time.Time      `gorm:"autoCreateTime" json:"created_at"`
	UpdatedAt     time.Time      `gorm:"autoUpdateTime" json:"updated_at"`
	ExpiresAt     *time.Time     `gorm:"index:IX_secrets_expires_at" json:"expires_at,omitempty"`
	DeletedAt     gorm.DeletedAt `gorm:"index" json:"-"`
}

func InitDB() (*DB, error) {
	dbPath := "onetime.db"

	db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
	if err != nil {
		return nil, err
	}

	// Get underlying sql.DB to configure connection pool if needed
	sqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}

	// Set connection pool settings
	sqlDB.SetMaxIdleConns(10)
	sqlDB.SetMaxOpenConns(100)
	sqlDB.SetConnMaxLifetime(time.Hour)

	return &DB{db}, nil
}

func MigrateDB(db *DB) error {
	return db.AutoMigrate(&Secret{})
}

func (db *DB) CreateSecret(secret *Secret) error {
	return db.Create(secret).Error
}

func (db *DB) GetSecretBySlug(slug string) (*Secret, error) {
	var secret Secret
	now := time.Now()

	err := db.Where("slug = ? AND (expires_at IS NULL OR expires_at > ?)", slug, now).First(&secret).Error
	if err != nil {
		return nil, err
	}

	return &secret, nil
}

func (db *DB) DeleteSecret(slug string) error {
	return db.Unscoped().Where("slug = ?", slug).Delete(&Secret{}).Error
}

func (db *DB) DeleteExpiredSecrets() error {
	now := time.Now()
	return db.Unscoped().Where("expires_at IS NOT NULL AND expires_at <= ?", now).Delete(&Secret{}).Error
}

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