aiforge/models/dataset.go

package models

import (
	"fmt"
	"sort"

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

const (
	DatasetStatusPrivate int = iota
	DatasetStatusPublic
	DatasetStatusDeleted
)

type DatasetList []*Dataset
type SearchDatasetOptions struct {
	ListOptions
	Keyword  string
	OwnerID  int64
	IsPublic bool
}

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"`
	CreatedUnix   timeutil.TimeStamp `xorm:"INDEX created"`
	UpdatedUnix   timeutil.TimeStamp `xorm:"INDEX updated"`

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

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

	return nil
}

// 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
}

// 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 UpdateDataset(ctx DBContext, rel *Dataset) error {
	_, err := ctx.e.ID(rel.ID).AllCols().Update(rel)
	return err
}

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.Eq{"user_id": opts.OwnerID})

	if opts.IsPublic {
		cond = cond.And(builder.Eq{"status": DatasetStatusPublic})
	} else {
		cond = cond.And(builder.Neq{"status": DatasetStatusDeleted})
	}

	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()

	// count, err := sess.Where(cond).Count(new(DatasetList))

	// if err != nil {
	// 	return nil, 0, fmt.Errorf("Count: %v", err)
	// }

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

	return repos, 0, 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 GeDatasetAttachments(rels ...*Dataset) (err error) {
	return geDatasetAttachments(x, rels...)
}

func geDatasetAttachments(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
}