brunel/db/repository.go

package db

import (
	"brunel/domain"
	"brunel/fastmap"
	"time"

	"github.com/jinzhu/now"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

const MaxBatchSize = 1999

type Repository struct {
	db *gorm.DB
}

func NewRepository(db *gorm.DB) *Repository {
	return &Repository{db: db}
}

func (r *Repository) GetUser(username string) (domain.User, error) {
	var user domain.User
	tx := r.db.Where("username = ?", username).First(&user)
	return user, tx.Error
}

func (r *Repository) CreateUser(user domain.User) error {
	tx := r.db.Create(&user)
	return tx.Error
}

func (r *Repository) UpdateUser(user domain.User) error {
	tx := r.db.Save(&user)
	return tx.Error
}

func (r *Repository) CreateSession(session domain.Session) error {
	tx := r.db.Create(&session)
	return tx.Error
}

func (r *Repository) GetSession(token string) (domain.Session, error) {
	now := now.BeginningOfMinute()
	var session domain.Session
	tx := r.db.Where("token = ? AND expiry > ?", token, now).First(&session)
	return session, tx.Error
}

func (r *Repository) DeleteSession(token string) error {
	tx := r.db.Delete(&domain.Session{}, "token = ?", token)
	return tx.Error
}

func (r *Repository) DeleteExpiredSessions() error {
	now := now.BeginningOfMinute()
	tx := r.db.Where("expiry < ?", now).Delete(&domain.Session{})
	return tx.Error
}

func (r *Repository) UpdatePackage(pkg domain.SourcePackage) error {
	dto := sourcePackageToDto(pkg)
	tx := r.db.Save(&dto)
	return tx.Error
}

func (r *Repository) GetPackages() ([]domain.SourcePackage, error) {
	var allPackages []domain.SourcePackage
	offset := 0

	for {
		var batchPackages []domain.SourcePackageDTO

		tx := r.db.Preload(clause.Associations).
			Offset(offset).
			Limit(MaxBatchSize).
			Find(&batchPackages)

		if tx.Error != nil {
			return nil, tx.Error
		}

		// Convert DTOs to domain objects
		for _, v := range batchPackages {
			allPackages = append(allPackages, sourcePackageDtoToDomain(v))
		}

		// If we've retrieved fewer records than MaxBatchSize, we're done
		if len(batchPackages) < MaxBatchSize {
			break
		}

		// Move to the next batch
		offset += MaxBatchSize
	}

	return allPackages, nil
}

func (r *Repository) GetPackage(name string) (domain.SourcePackage, error) {
	var pkg domain.SourcePackageDTO
	tx := r.db.Where("name = ?", name).First(&pkg)
	if tx.Error != nil {
		return domain.SourcePackage{}, tx.Error
	}
	return sourcePackageDtoToDomain(pkg), nil
}

func (r *Repository) SavePackages(pkgs *fastmap.Fastmap[string, domain.SourcePackage]) error {
	packs := make([]domain.SourcePackageDTO, 0)
	pkgs.Iter(func(k string, v domain.SourcePackage) bool {
		packs = append(packs, sourcePackageToDto(v))
		return true
	})

	for i := 0; i < len(packs); i += MaxBatchSize {
		end := i + MaxBatchSize
		length := len(packs)
		if end > length {
			end = length
		}
		tx := r.db.Save(packs[i:end])
		if tx.Error != nil {
			return tx.Error
		}
	}

	return nil
}

func (r *Repository) UpdateLastUpdateTime(time time.Time) error {
	val := &domain.TimeContainer{
		Time: time,
		ID:   "lastupdatetime",
	}
	tx := r.db.Save(val)
	return tx.Error
}

func sourcePackageToDto(pkg domain.SourcePackage) domain.SourcePackageDTO {
	dto := domain.SourcePackageDTO{
		Name:     pkg.Name,
		Packages: make([]domain.PackageInfo, 0),
	}
	pkg.Packages.Iter(func(k string, v domain.PackageInfo) bool {
		dto.Packages = append(dto.Packages, v)
		return true
	})
	return dto
}

func sourcePackageDtoToDomain(dto domain.SourcePackageDTO) domain.SourcePackage {
	pkg := domain.SourcePackage{
		Name:     dto.Name,
		Packages: fastmap.New[string, domain.PackageInfo](),
	}
	for _, v := range dto.Packages {
		pkg.Packages.Set(v.PackageName, v)
	}
	return pkg
}