aiforge/models/dataset.go

package models

import (
	"fmt"

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