aiforge/models/dataset.go

package models

import (
	"errors"
	"fmt"
	"sort"

	"code.gitea.io/gitea/modules/timeutil"
	"xorm.io/builder"
)

const (
	DatasetStatusPrivate int32 = iota
	DatasetStatusPublic
	DatasetStatusDeleted
)

type Dataset struct {
	ID            int64  `xorm:"pk autoincr"`
	Title         string `xorm:"INDEX NOT NULL"`
	Status        int32  `xorm:"INDEX"` // normal_private: 0, pulbic: 1, is_delete: 2
	Category      string
	Description   string `xorm:"TEXT"`
	DownloadTimes int64
	License       string
	Task          string
	ReleaseID     int64              `xorm:"INDEX"`
	UserID        int64              `xorm:"INDEX"`
	RepoID        int64              `xorm:"INDEX"`
	CreatedUnix   timeutil.TimeStamp `xorm:"INDEX created"`
	UpdatedUnix   timeutil.TimeStamp `xorm:"INDEX updated"`

	User        *User         `xorm:"-"`
	Attachments []*Attachment `xorm:"-"`
}

func (d *Dataset) IsPrivate() bool {
	switch d.Status {
	case DatasetStatusPrivate:
		return true
	case DatasetStatusPublic:
		return false
	case DatasetStatusDeleted:
		return false
	default:
		return false
	}
}

type DatasetList []*Dataset

func (datasets DatasetList) loadAttributes(e Engine) error {
	if len(datasets) == 0 {
		return nil
	}

	set := make(map[int64]struct{})
	datasetIDs := make([]int64, len(datasets))
	for i := range datasets {
		set[datasets[i].UserID] = struct{}{}
		datasetIDs[i] = datasets[i].ID
	}

	// Load owners.
	users := make(map[int64]*User, len(set))
	if err := e.
		Where("id > 0").
		In("id", keysInt64(set)).
		Find(&users); err != nil {
		return fmt.Errorf("find users: %v", err)
	}
	for i := range datasets {
		datasets[i].User = users[datasets[i].UserID]
	}

	return nil
}

type SearchDatasetOptions struct {
	Keyword       string
	OwnerID       int64
	RepoID        int64
	IncludePublic bool
	ListOptions
	SearchOrderBy
}

func CreateDataset(dataset *Dataset) (err error) {
	if _, err = x.Insert(dataset); err != nil {
		return err
	}

	return nil
}

func SearchDataset(opts *SearchDatasetOptions) (DatasetList, int64, error) {
	cond := SearchDatasetCondition(opts)
	return SearchDatasetByCondition(opts, cond)
}

func SearchDatasetCondition(opts *SearchDatasetOptions) builder.Cond {
	var cond = builder.NewCond()
	cond = cond.And(builder.Neq{"status": DatasetStatusDeleted})

	if len(opts.Keyword) > 0 {
		cond = cond.And(builder.Like{"title", opts.Keyword})
	}

	if opts.RepoID > 0 {
		cond = cond.And(builder.Eq{"repo_id": opts.RepoID})
	}

	if opts.IncludePublic {
		cond = cond.And(builder.Eq{"status": DatasetStatusPublic})
		if opts.OwnerID > 0 {
			cond = cond.Or(builder.Eq{"user_id": opts.OwnerID})
		}
	} else {
		if opts.OwnerID > 0 {
			cond = cond.And(builder.Eq{"user_id": opts.OwnerID})
		}
	}

	return cond
}

func SearchDatasetByCondition(opts *SearchDatasetOptions, cond builder.Cond) (DatasetList, int64, error) {
	if opts.Page <= 0 {
		opts.Page = 1
	}

	var err error
	sess := x.NewSession()
	defer sess.Close()

	datasets := make(DatasetList, 0, opts.PageSize)

	count, err := sess.Where(cond).Count(new(Dataset))
	if err != nil {
		return nil, 0, fmt.Errorf("Count: %v", err)
	}

	sess.Where(cond).OrderBy(opts.SearchOrderBy.String())
	if opts.PageSize > 0 {
		sess.Limit(opts.PageSize, (opts.Page-1)*opts.PageSize)
	}
	if err = sess.Find(&datasets); err != nil {
		return nil, 0, fmt.Errorf("Dataset: %v", err)
	}

	if err = datasets.loadAttributes(sess); err != nil {
		return nil, 0, fmt.Errorf("LoadAttributes: %v", err)
	}

	return datasets, count, nil
}

type datasetMetaSearch struct {
	ID  []int64
	Rel []*Dataset
}

func (s datasetMetaSearch) Len() int {
	return len(s.ID)
}
func (s datasetMetaSearch) Swap(i, j int) {
	s.ID[i], s.ID[j] = s.ID[j], s.ID[i]
	s.Rel[i], s.Rel[j] = s.Rel[j], s.Rel[i]
}
func (s datasetMetaSearch) Less(i, j int) bool {
	return s.ID[i] < s.ID[j]
}

func GetDatasetAttachments(rels ...*Dataset) (err error) {
	return getDatasetAttachments(x, rels...)
}

func getDatasetAttachments(e Engine, rels ...*Dataset) (err error) {
	if len(rels) == 0 {
		return
	}

	// To keep this efficient as possible sort all datasets by id,
	//    select attachments by dataset id,
	//    then merge join them

	// Sort
	var sortedRels = datasetMetaSearch{ID: make([]int64, len(rels)), Rel: make([]*Dataset, len(rels))}
	var attachments []*Attachment
	for index, element := range rels {
		element.Attachments = []*Attachment{}
		sortedRels.ID[index] = element.ID
		sortedRels.Rel[index] = element
	}
	sort.Sort(sortedRels)

	// Select attachments
	err = e.
		Asc("dataset_id").
		In("dataset_id", sortedRels.ID).
		Find(&attachments, Attachment{})
	if err != nil {
		return err
	}

	// merge join
	var currentIndex = 0
	for _, attachment := range attachments {
		for sortedRels.ID[currentIndex] < attachment.DatasetID {
			currentIndex++
		}
		sortedRels.Rel[currentIndex].Attachments = append(sortedRels.Rel[currentIndex].Attachments, attachment)
	}

	return
}

// AddDatasetAttachments adds a Dataset attachments
func AddDatasetAttachments(DatasetID int64, attachmentUUIDs []string) (err error) {
	// Check attachments
	attachments, err := GetAttachmentsByUUIDs(attachmentUUIDs)
	if err != nil {
		return fmt.Errorf("GetAttachmentsByUUIDs [uuids: %v]: %v", attachmentUUIDs, err)
	}

	for i := range attachments {
		attachments[i].DatasetID = DatasetID
		// No assign value could be 0, so ignore AllCols().
		if _, err = x.ID(attachments[i].ID).Update(attachments[i]); err != nil {
			return fmt.Errorf("update attachment [%d]: %v", attachments[i].ID, err)
		}
	}

	return
}

func UpdateDataset(ctx DBContext, rel *Dataset) error {
	_, err := ctx.e.ID(rel.ID).AllCols().Update(rel)
	return err
}

// GetDatasetByID returns Dataset with given ID.
func GetDatasetByID(id int64) (*Dataset, error) {
	rel := new(Dataset)
	has, err := x.
		ID(id).
		Get(rel)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrDatasetNotExist{id}
	}

	return rel, nil
}

func DeleteDataset(datasetID int64, uid int64) error {
	var err error
	sess := x.NewSession()
	defer sess.Close()
	if err = sess.Begin(); err != nil {
		return err
	}

	dataset := &Dataset{ID: datasetID, UserID: uid}
	has, err := sess.Get(dataset)
	if err != nil {
		return err
	} else if !has {
		return errors.New("not found")
	}

	if cnt, err := sess.ID(datasetID).Delete(new(Dataset)); err != nil {
		return err
	} else if cnt != 1 {
		return errors.New("not found")
	}
	if err = sess.Commit(); err != nil {
		sess.Close()
		return fmt.Errorf("Commit: %v", err)
	}
	return nil
}

func GetOwnerDatasetByID(id int64, user *User) (*Dataset, error) {
	dataset, err := GetDatasetByID(id)
	if err != nil {
		return nil, err
	}

	if !dataset.IsPrivate() {
		return dataset, nil
	}
	if dataset.IsPrivate() && user != nil && user.ID == dataset.UserID {
		return dataset, nil
	}
	return nil, errors.New("dataset not fount")
}

func IncreaseDownloadCount(datasetID int64) error {
	// Update download count.
	if _, err := x.Exec("UPDATE `dataset` SET download_times=download_times+1 WHERE id=?", datasetID); err != nil {
		return fmt.Errorf("increase dataset count: %v", err)
	}

	return nil
}