aiforge/routers/repo/dataset.go

package repo

import (
	"encoding/json"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"unicode/utf8"

	"code.gitea.io/gitea/services/repository"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/auth"
	"code.gitea.io/gitea/modules/base"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
)

const (
	tplIndex         base.TplName = "repo/datasets/index"
	tplDatasetCreate base.TplName = "repo/datasets/create"
	tplDatasetEdit   base.TplName = "repo/datasets/edit"
	taskstplIndex    base.TplName = "repo/datasets/tasks/index"
	tplReference     base.TplName = "repo/datasets/reference"
)

// MustEnableDataset check if repository enable internal dataset
func MustEnableDataset(ctx *context.Context) {
	if !ctx.Repo.CanRead(models.UnitTypeDatasets) {
		ctx.NotFound("MustEnableDataset", nil)
		return
	}
}

func newFilterPrivateAttachments(ctx *context.Context, list []*models.Attachment, repo *models.Repository) []*models.Attachment {

	if ctx.Repo.CanWrite(models.UnitTypeDatasets) {
		log.Info("can write.")
		return list
	} else {
		if repo.Owner == nil {
			repo.GetOwner()
		}
		permission := false
		if !permission && ctx.User != nil {
			isCollaborator, _ := repo.IsCollaborator(ctx.User.ID)
			isInRepoTeam, _ := repo.IsInRepoTeam(ctx.User.ID)
			if isCollaborator || isInRepoTeam {
				log.Info("Collaborator user may visit the attach.")
				permission = true
			}
		}

		var publicList []*models.Attachment
		for _, attach := range list {
			if !attach.IsPrivate {
				publicList = append(publicList, attach)
			} else {
				if permission {
					publicList = append(publicList, attach)
				}
			}
		}
		return publicList
	}
}

func QueryDataSet(ctx *context.Context) []*models.Attachment {
	repo := ctx.Repo.Repository

	dataset, err := models.GetDatasetByRepo(repo)
	if err != nil {
		log.Error("zou  not found dataset 1")
		ctx.NotFound("GetDatasetByRepo", err)
		return nil
	}

	if ctx.Query("type") == "" {
		log.Error("zou  not found  type 2")
		ctx.NotFound("type error", nil)
		return nil
	}
	err = models.GetDatasetAttachments(ctx.QueryInt("type"), ctx.IsSigned, ctx.User, dataset)
	if err != nil {
		ctx.ServerError("GetDatasetAttachments", err)
		return nil
	}
	attachments := newFilterPrivateAttachments(ctx, dataset.Attachments, repo)

	ctx.Data["SortType"] = ctx.Query("sort")

	sort.Slice(attachments, func(i, j int) bool {
		return attachments[i].CreatedUnix > attachments[j].CreatedUnix
	})

	return attachments
}

func DatasetIndex(ctx *context.Context) {
	log.Info("dataset index 1")
	MustEnableDataset(ctx)
	ctx.Data["PageIsDataset"] = true

	ctx.Data["SortType"] = ctx.Query("sort")

	repo := ctx.Repo.Repository

	dataset, err := models.GetDatasetByRepo(repo)
	ctx.Data["CanWrite"] = ctx.Repo.CanWrite(models.UnitTypeDatasets)
	if err != nil {
		log.Warn("query dataset, not found.")
		ctx.HTML(200, tplIndex)
		return
	}
	cloudbrainType := -1
	if ctx.Query("type") != "" {

		cloudbrainType = ctx.QueryInt("type")
	}
	err = models.GetDatasetAttachments(cloudbrainType, ctx.IsSigned, ctx.User, dataset)
	if err != nil {
		ctx.ServerError("GetDatasetAttachments", err)
		return
	}

	attachments := newFilterPrivateAttachments(ctx, dataset.Attachments, repo)

	if ctx.Data["SortType"] == "nameAsc" {
		sort.Slice(attachments, func(i, j int) bool {
			return strings.ToLower(attachments[i].Name) < strings.ToLower(attachments[j].Name)
		})
	} else if ctx.Data["SortType"] == "nameDesc" {
		sort.Slice(attachments, func(i, j int) bool {
			return strings.ToLower(attachments[i].Name) > strings.ToLower(attachments[j].Name)
		})
	} else if ctx.Data["SortType"] == "sizeAsc" {
		sort.Slice(attachments, func(i, j int) bool {
			return attachments[i].Size < attachments[j].Size
		})
	} else if ctx.Data["SortType"] == "sizeDesc" {
		sort.Slice(attachments, func(i, j int) bool {
			return attachments[i].Size > attachments[j].Size
		})
	} else if ctx.Data["SortType"] == "timeAsc" {
		sort.Slice(attachments, func(i, j int) bool {
			return attachments[i].CreatedUnix < attachments[j].CreatedUnix
		})
	} else {
		sort.Slice(attachments, func(i, j int) bool {
			return attachments[i].CreatedUnix > attachments[j].CreatedUnix
		})
	}

	page := ctx.QueryInt("page")
	if page <= 0 {
		page = 1
	}
	pagesize := ctx.QueryInt("pagesize")
	if pagesize <= 0 {
		pagesize = 10
	}
	pager := context.NewPagination(len(attachments), pagesize, page, 5)

	pageAttachments := getPageAttachments(attachments, page, pagesize)

	//load attachment creator
	for _, attachment := range pageAttachments {
		uploader, _ := models.GetUserByID(attachment.UploaderID)
		attachment.Uploader = uploader
		if !strings.HasSuffix(attachment.Name, ".zip") && !strings.HasSuffix(attachment.Name, ".tar.gz") {
			attachment.DecompressState = -1 //非压缩文件
		}

	}

	ctx.Data["Page"] = pager

	ctx.Data["Title"] = ctx.Tr("dataset.show_dataset")
	ctx.Data["Link"] = ctx.Repo.RepoLink + "/datasets"
	ctx.Data["dataset"] = dataset
	ctx.Data["Attachments"] = pageAttachments
	ctx.Data["IsOwner"] = true
	ctx.Data["StoreType"] = setting.Attachment.StoreType
	ctx.Data["Type"] = cloudbrainType

	renderAttachmentSettings(ctx)

	ctx.HTML(200, tplIndex)
}

func getPageAttachments(attachments []*models.Attachment, page int, pagesize int) []*models.Attachment {
	begin := (page - 1) * pagesize
	end := (page) * pagesize

	if begin > len(attachments)-1 {
		return nil
	}
	if end > len(attachments)-1 {
		return attachments[begin:]
	} else {
		return attachments[begin:end]
	}

}

func CreateDataset(ctx *context.Context) {

	MustEnableDataset(ctx)
	ctx.Data["PageIsDataset"] = true

	ctx.HTML(200, tplDatasetCreate)
}

func EditDataset(ctx *context.Context) {

	MustEnableDataset(ctx)
	ctx.Data["PageIsDataset"] = true
	datasetId, _ := strconv.ParseInt(ctx.Params(":id"), 10, 64)

	dataset, _ := models.GetDatasetByID(datasetId)
	if dataset == nil {
		ctx.Error(http.StatusNotFound, "")
		return
	}
	ctx.Data["Dataset"] = dataset

	ctx.HTML(200, tplDatasetEdit)
}

func CreateDatasetPost(ctx *context.Context, form auth.CreateDatasetForm) {

	dataset := &models.Dataset{}

	if !NamePattern.MatchString(form.Title) {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("dataset.title_format_err")))
		return
	}
	if utf8.RuneCountInString(form.Description) > 1024 {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("dataset.description_format_err", 1024)))
		return
	}

	dataset.RepoID = ctx.Repo.Repository.ID
	dataset.UserID = ctx.User.ID
	dataset.Category = form.Category
	dataset.Task = form.Task
	dataset.Title = form.Title
	dataset.License = form.License
	dataset.Description = form.Description
	dataset.DownloadTimes = 0
	if ctx.Repo.Repository.IsPrivate {
		dataset.Status = 0
	} else {
		dataset.Status = 1
	}
	err := models.CreateDataset(dataset)
	if err != nil {
		log.Error("fail to create dataset", err)
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("dataset.create_dataset_fail")))
	} else {
		ctx.JSON(http.StatusOK, models.BaseOKMessage)
	}

}
func ReferenceDatasetDelete(ctx *context.Context) {
	repoID := ctx.Repo.Repository.ID
	datasetId, _ := strconv.ParseInt(ctx.Params(":id"), 10, 64)

	oldDatasetIds := models.GetDatasetIdsByRepoID(repoID)

	var newDatasetIds []int64

	for _, tempDatasetId := range oldDatasetIds {
		if datasetId != tempDatasetId {
			newDatasetIds = append(newDatasetIds, tempDatasetId)
		}
	}
	err := models.NewDatasetIdsByRepoID(repoID, newDatasetIds)
	if err != nil {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage("dataset.cancel_reference_dataset_fail"))
	}
	ctx.JSON(http.StatusOK, models.BaseOKMessage)

}

func ReferenceDatasetPost(ctx *context.Context, form auth.ReferenceDatasetForm) {
	repoID := ctx.Repo.Repository.ID
	err := models.NewDatasetIdsByRepoID(repoID, form.DatasetID)
	if err != nil {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage("dataset.reference_dataset_fail"))
	}

	ctx.JSON(http.StatusOK, models.BaseOKMessage)

}

func EditDatasetPost(ctx *context.Context, form auth.EditDatasetForm) {
	ctx.Data["PageIsDataset"] = true

	ctx.Data["Title"] = ctx.Tr("dataset.edit_dataset")

	if !NamePattern.MatchString(form.Title) {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("dataset.title_format_err")))
		return
	}
	if utf8.RuneCountInString(form.Description) > 1024 {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("dataset.description_format_err", 1024)))
		return
	}

	rel, err := models.GetDatasetByID(form.ID)
	ctx.Data["dataset"] = rel

	if err != nil {
		log.Error("failed to query dataset", err)
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("dataset.query_dataset_fail")))
		return
	}

	rel.Title = form.Title
	rel.Description = form.Description
	rel.Category = form.Category
	rel.Task = form.Task
	rel.License = form.License
	if err = models.UpdateDataset(models.DefaultDBContext(), rel); err != nil {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("dataset.query_dataset_fail")))
	}
	ctx.JSON(http.StatusOK, models.BaseOKMessage)
}

func DatasetAction(ctx *context.Context) {
	var err error
	datasetId, _ := strconv.ParseInt(ctx.Params(":id"), 10, 64)
	switch ctx.Params(":action") {
	case "star":
		err = models.StarDataset(ctx.User.ID, datasetId, true)
	case "unstar":
		err = models.StarDataset(ctx.User.ID, datasetId, false)

	}
	if err != nil {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr("repo.star_fail", ctx.Params(":action"))))
	} else {
		ctx.JSON(http.StatusOK, models.BaseOKMessage)
	}

}

func datasetMultiple(ctx *context.Context, opts *models.SearchDatasetOptions) {
	page := ctx.QueryInt("page")
	keyword := strings.Trim(ctx.Query("q"), " ")
	opts.Keyword = keyword
	if opts.SearchOrderBy.String() == "" {
		opts.SearchOrderBy = models.SearchOrderByRecentUpdated
	}

	opts.RecommendOnly = ctx.QueryBool("recommend")
	opts.ListOptions = models.ListOptions{
		Page:     page,
		PageSize: setting.UI.DatasetPagingNum,
	}
	opts.JustNeedZipFile = true
	opts.User = ctx.User

	datasets, count, err := models.SearchDataset(opts)

	if err != nil {
		ctx.ServerError("datasets", err)
		return
	}

	data, err := json.Marshal(datasets)
	log.Info("datakey", string(data))
	if err != nil {
		log.Error("json.Marshal failed:", err.Error())
		ctx.JSON(200, map[string]string{
			"result_code": "-1",
			"error_msg":   err.Error(),
			"data":        "",
		})
		return
	}
	ctx.JSON(200, map[string]string{
		"result_code": "0",
		"data":        string(data),
		"count":       strconv.FormatInt(count, 10),
	})
}

func CurrentRepoDatasetMultiple(ctx *context.Context) {
	datasetIds := models.GetDatasetIdsByRepoID(ctx.Repo.Repository.ID)
	searchOrderBy := getSearchOrderByInValues(datasetIds)
	opts := &models.SearchDatasetOptions{
		RepoID:         ctx.Repo.Repository.ID,
		NeedAttachment: true,
		CloudBrainType: ctx.QueryInt("type"),
		DatasetIDs:     datasetIds,
		SearchOrderBy:  searchOrderBy,
	}

	datasetMultiple(ctx, opts)

}

func getSearchOrderByInValues(datasetIds []int64) models.SearchOrderBy {
	if len(datasetIds) == 0 {
		return ""
	}
	searchOrderBy := "CASE id "
	for i, id := range datasetIds {
		searchOrderBy += fmt.Sprintf(" WHEN %d THEN %d", id, i+1)
	}
	searchOrderBy += " ELSE 0 END"
	return models.SearchOrderBy(searchOrderBy)
}

func MyDatasetsMultiple(ctx *context.Context) {

	opts := &models.SearchDatasetOptions{
		UploadAttachmentByMe: true,
		NeedAttachment:       true,
		CloudBrainType:       ctx.QueryInt("type"),
	}
	datasetMultiple(ctx, opts)

}

func ReferenceDatasetAvailable(ctx *context.Context) {

	opts := &models.SearchDatasetOptions{
		PublicOnly:     true,
		NeedAttachment: false,
		CloudBrainType: models.TypeCloudBrainAll,
		SearchOrderBy:  models.SearchOrderByDefault,
	}
	dataset, _ := models.GetDatasetByRepo(&models.Repository{ID: ctx.Repo.Repository.ID})
	if dataset != nil {
		opts.ExcludeDatasetId = dataset.ID
	}
	datasetMultiple(ctx, opts)

}

func PublicDatasetMultiple(ctx *context.Context) {

	opts := &models.SearchDatasetOptions{
		PublicOnly:     true,
		NeedAttachment: true,
		CloudBrainType: ctx.QueryInt("type"),
		SearchOrderBy:  models.SearchOrderByDefault,
	}
	datasetMultiple(ctx, opts)

}

func MyFavoriteDatasetMultiple(ctx *context.Context) {

	opts := &models.SearchDatasetOptions{
		StarByMe:       true,
		DatasetIDs:     models.GetDatasetIdsStarByUser(ctx.User.ID),
		NeedAttachment: true,
		CloudBrainType: ctx.QueryInt("type"),
	}
	datasetMultiple(ctx, opts)
}
func ReferenceDataset(ctx *context.Context) {
	MustEnableDataset(ctx)
	ctx.Data["PageIsDataset"] = true
	ctx.Data["MaxReferenceDatasetNum"] = setting.RepoMaxReferenceDatasetNum
	ctx.Data["CanWrite"] = ctx.Repo.CanWrite(models.UnitTypeDatasets)
	ctx.HTML(200, tplReference)

}

func ReferenceDatasetData(ctx *context.Context) {
	MustEnableDataset(ctx)
	datasetIds := models.GetDatasetIdsByRepoID(ctx.Repo.Repository.ID)
	var datasets models.DatasetList
	var err error
	if len(datasetIds) > 0 {

		opts := &models.SearchDatasetOptions{
			DatasetIDs:     datasetIds,
			NeedAttachment: false,
			CloudBrainType: models.TypeCloudBrainAll,
			ListOptions: models.ListOptions{
				Page:     1,
				PageSize: setting.RepoMaxReferenceDatasetNum,
			},
			SearchOrderBy:  getSearchOrderByInValues(datasetIds),
			QueryReference: true,
		}
		datasets, _, err = models.SearchDataset(opts)
		if err != nil {
			ctx.ServerError("SearchDatasets", err)
			return
		}
	}

	ctx.JSON(http.StatusOK, repository.ConvertToDatasetWithStar(ctx, datasets))

}

func GetDatasetStatus(ctx *context.Context) {

	var (
		err error
	)

	UUID := ctx.Params(":uuid")
	attachment, err := models.GetAttachmentByUUID(UUID)
	if err != nil {
		log.Error("GetDatasetStarByUser failed:", err.Error())
		ctx.JSON(200, map[string]string{
			"result_code": "-1",
			"error_msg":   err.Error(),
			"data":        "",
		})
		return
	}

	ctx.JSON(200, map[string]string{
		"result_code":      "0",
		"UUID":             UUID,
		"AttachmentStatus": fmt.Sprint(attachment.DecompressState),
	})
}