aiforge/routers/search.go

package routers

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"github.com/olivere/elastic/v7"
)

type SearchRes struct {
	Total  int64
	Result []map[string]interface{}
}

var client *elastic.Client

func InitESClient() {
	ESSearchUrl := setting.ESSearchURL
	var err error
	client, err = elastic.NewClient(elastic.SetSniff(false), elastic.SetURL(ESSearchUrl))
	if err != nil {
		panic(err)
	}
}

func Search(ctx *context.Context) {
	TableName := ctx.Query("TableName")
	Key := ctx.Query("Key")
	Page := ctx.QueryInt("Page")
	PageSize := ctx.QueryInt("PageSize")
	if Page <= 0 {
		Page = 1
	}
	if PageSize <= 0 || PageSize > 200 {
		PageSize = setting.UI.IssuePagingNum
	}

	if TableName == "repository" {
		searchRepo(ctx, "repository-es-index", Key, Page, PageSize)
		return
	} else if TableName == "issue" {
		searchIssue(ctx, "issue-es-index", Key, Page, PageSize)
		return
	} else if TableName == "user" {
		searchUserOrOrg(ctx, "user-es-index", Key, Page, PageSize, true)
		return
	} else if TableName == "org" {
		searchUserOrOrg(ctx, "user-es-index", Key, Page, PageSize, false)
		return
	} else if TableName == "dataset" {
		searchDataSet(ctx, "dataset-es-index", Key, Page, PageSize)
		return
	} else if TableName == "pr" {
		searchPR(ctx, "issue-es-index", Key, Page, PageSize)
		return
	}
}

func searchRepoByLabel(ctx *context.Context, TableName string, Key string, Page int, PageSize int) {
	/*
			   项目， ES名称： repository-es-index
			   搜索:
			   name character varying(255) ,  项目名称
			   description text,  项目描述
			   topics json,  标签
			   排序:
			   updated_unix
			   num_watches,
		       num_stars,
		       num_forks,
	*/
	SortBy := ctx.Query("SortBy")
	if SortBy == "" {
		SortBy = "updated_unix.keyword"
	}
	ascending := ctx.QueryBool("Ascending")
	log.Info("query searchRepoByLabel start")
	if Key != "" {
		boolQ := elastic.NewBoolQuery()
		topicsQuery := elastic.NewMatchQuery("topics", Key)
		boolQ.Should(topicsQuery)
		res, err := client.Search(TableName).Query(boolQ).Sort(SortBy, ascending).From((Page - 1) * PageSize).Size(PageSize).Do(ctx.Req.Context())
		if err == nil {
			result := makeRepoResult(res, "")
			ctx.JSON(200, result)
			return
		} else {
			log.Info("query es error," + err.Error())
		}
	}
	ctx.JSON(200, "")
}

func searchRepo(ctx *context.Context, TableName string, Key string, Page int, PageSize int) {
	/*
			   项目， ES名称： repository-es-index
			   搜索:
			   name character varying(255) ,  项目名称
			   description text,  项目描述
			   topics json,  标签
			   排序:
			   updated_unix
			   num_watches,
		       num_stars,
		       num_forks,
	*/

	SortBy := ctx.Query("SortBy")
	if SortBy == "" {
		SortBy = "updated_unix.keyword"
	}
	ascending := ctx.QueryBool("Ascending")
	log.Info("query searchRepo start")
	if Key != "" {
		boolQ := elastic.NewBoolQuery()
		nameQuery := elastic.NewMatchQuery("name", Key).Boost(1024).QueryName("f_first")
		descriptionQuery := elastic.NewMatchQuery("description", Key).Boost(1.5).QueryName("f_second")
		topicsQuery := elastic.NewMatchQuery("topics", Key).Boost(1).QueryName("f_third")
		boolQ.Should(nameQuery, descriptionQuery, topicsQuery)
		res, err := client.Search(TableName).Query(boolQ).SortBy(elastic.NewScoreSort(), elastic.NewFieldSort(SortBy).Order(ascending)).From((Page - 1) * PageSize).Size(PageSize).Do(ctx.Req.Context())
		if err == nil {
			result := makeRepoResult(res, Key)
			ctx.JSON(200, result)
		} else {
			log.Info("query es error," + err.Error())
			ctx.JSON(200, "")
		}
	} else {
		log.Info("query all content.")
		//搜索的属性要指定{"timestamp":{"unmapped_type":"date"}}
		res, err := client.Search(TableName).SortBy(elastic.NewFieldSort(SortBy).Order(ascending)).From((Page - 1) * PageSize).Size(PageSize).Do(ctx.Req.Context())
		if err == nil {
			result := makeRepoResult(res, "")
			ctx.JSON(200, result)
		} else {
			log.Info("query es error," + err.Error())
			ctx.JSON(200, "")
		}
	}
}

func makeRepoResult(sRes *elastic.SearchResult, Key string) *SearchRes {
	total := sRes.Hits.TotalHits.Value
	result := make([]map[string]interface{}, 0)

	for i, hit := range sRes.Hits.Hits {
		log.Info("this is repo query " + fmt.Sprint(i) + " result.")
		recordSource := make(map[string]interface{})
		source, err := hit.Source.MarshalJSON()

		if err == nil {
			err = json.Unmarshal(source, &recordSource)
			if err == nil {
				record := make(map[string]interface{})
				record["name"] = recordSource["name"]
				record["owner_name"] = recordSource["owner_name"]
				desc := recordSource["description"].(string)

				record["description"] = dealLongText(desc, Key, hit.MatchedQueries)

				record["num_watches"] = recordSource["num_watches"]
				record["num_stars"] = recordSource["num_stars"]
				record["num_forks"] = recordSource["num_forks"]
				record["topics"] = recordSource["topics"]
				if recordSource["avatar"] != nil {
					record["avatar"] = setting.AppSubURL + "/repo-avatars/" + recordSource["avatar"].(string)
				}
				record["updated_unix"] = recordSource["updated_unix"]
				record["lang"] = recordSource["lang"]
				result = append(result, record)
			} else {
				log.Info("deal repo source error," + err.Error())
			}
		} else {
			log.Info("deal repo source error," + err.Error())
		}
	}

	returnObj := &SearchRes{
		Total:  total,
		Result: result,
	}

	return returnObj
}

func dealLongText(text string, Key string, MatchedQueries []string) string {
	var isNeedToDealText bool
	isNeedToDealText = false
	if len(MatchedQueries) > 0 && Key != "" {
		if MatchedQueries[0] == "f_second" || MatchedQueries[0] == "f_third" {
			isNeedToDealText = true
		}
	}
	stringlen := len(text)
	if isNeedToDealText && stringlen > 200 {
		index := strings.Index(text, Key)
		if index > 0 {
			start := index - 50
			if start < 0 {
				start = 0
			}
			end := index + 150
			if end >= stringlen {
				end = stringlen
			}
			return text[start:end]
		} else {
			return text[0:200]
		}
	} else {
		if stringlen > 200 {
			return text[0:200]
		} else {
			return text
		}
	}

}

func searchUserOrOrg(ctx *context.Context, TableName string, Key string, Page int, PageSize int, IsQueryUser bool) {
	/*
	   用户或者组织 ES名称： user-es-index
	   搜索:
	   name  ,  名称
	   full_name   全名
	   description  描述或者简介
	   排序:
	   created_unix
	   名称字母序
	*/
	SortBy := ctx.Query("SortBy")
	if SortBy == "" {
		SortBy = "updated_unix.keyword"
	}
	ascending := ctx.QueryBool("Ascending")
	boolQ := elastic.NewBoolQuery()
	if Key != "" {
		nameQuery := elastic.NewMatchQuery("name", Key).Boost(2).QueryName("f_first")
		full_nameQuery := elastic.NewMatchQuery("full_name", Key).Boost(1.5).QueryName("f_second")
		descriptionQuery := elastic.NewMatchQuery("description", Key).Boost(1).QueryName("f_third")
		boolQ.Should(nameQuery, full_nameQuery, descriptionQuery)
	}
	typeValue := 1
	if IsQueryUser {
		typeValue = 0
	}
	UserOrOrgQuery := elastic.NewTermQuery("type", typeValue)
	boolQ.Must(UserOrOrgQuery)

	res, err := client.Search(TableName).Query(boolQ).Sort(SortBy, ascending).From((Page - 1) * PageSize).Size(PageSize).Do(ctx.Req.Context())
	if err == nil {
		result := makeUserOrOrgResult(res, Key, ctx)
		ctx.JSON(200, result)
	} else {
		log.Info("query es error," + err.Error())
		ctx.JSON(200, "")
	}
}

func makeUserOrOrgResult(sRes *elastic.SearchResult, Key string, ctx *context.Context) *SearchRes {
	total := sRes.Hits.TotalHits.Value
	result := make([]map[string]interface{}, 0)

	for i, hit := range sRes.Hits.Hits {
		log.Info("this is user query " + fmt.Sprint(i) + " result.")
		recordSource := make(map[string]interface{})
		source, err := hit.Source.MarshalJSON()

		if err == nil {
			err = json.Unmarshal(source, &recordSource)
			if err == nil {
				record := make(map[string]interface{})
				record["name"] = recordSource["name"]
				record["full_name"] = recordSource["full_name"]
				desc := recordSource["description"].(string)

				record["description"] = dealLongText(desc, Key, hit.MatchedQueries)
				if ctx.User != nil {
					record["email"] = recordSource["email"]
				} else {
					record["email"] = ""
				}

				record["location"] = recordSource["location"]
				record["website"] = recordSource["website"]
				record["num_repos"] = recordSource["num_repos"]
				record["num_teams"] = recordSource["num_teams"]
				record["num_members"] = recordSource["num_members"]

				record["avatar"] = strings.TrimRight(setting.AppSubURL, "/") + "/user/avatar/" + recordSource["name"].(string) + "/" + strconv.Itoa(-1)
				record["updated_unix"] = recordSource["updated_unix"]
				record["created_unix"] = recordSource["created_unix"]
				result = append(result, record)
			} else {
				log.Info("deal user source error," + err.Error())
			}
		} else {
			log.Info("deal user source error," + err.Error())
		}
	}

	returnObj := &SearchRes{
		Total:  total,
		Result: result,
	}

	return returnObj
}

func searchDataSet(ctx *context.Context, TableName string, Key string, Page int, PageSize int) {
	/*
	   数据集，ES名称：dataset-es-index
	   搜索:
	   title  ,  名称
	   description   描述
	   category  标签
	   file_name  数据集文件名称
	   排序:
	   download_times

	*/
	SortBy := ctx.Query("SortBy")
	if SortBy == "" {
		SortBy = "download_times.keyword"
	}
	ascending := ctx.QueryBool("Ascending")
	log.Info("query searchRepo start")
	boolQ := elastic.NewBoolQuery()
	if Key != "" {
		nameQuery := elastic.NewMatchQuery("title", Key).Boost(2).QueryName("f_first")
		descQuery := elastic.NewMatchQuery("description", Key).Boost(1.5).QueryName("f_second")
		fileNameQuery := elastic.NewMatchQuery("file_name", Key).Boost(1).QueryName("f_third")
		categoryQuery := elastic.NewMatchQuery("category", Key).Boost(1).QueryName("f_fourth")
		boolQ.Should(nameQuery, descQuery, categoryQuery, fileNameQuery)
		res, err := client.Search(TableName).Query(boolQ).Sort(SortBy, ascending).From((Page - 1) * PageSize).Size(PageSize).Do(ctx.Req.Context())
		if err == nil {
			result := makeDatasetResult(res, Key)
			ctx.JSON(200, result)
		} else {
			log.Info("query es error," + err.Error())
		}
	} else {
		log.Info("query all content.")
		//搜索的属性要指定{"timestamp":{"unmapped_type":"date"}}
		res, err := client.Search(TableName).Sort(SortBy, ascending).From((Page - 1) * PageSize).Size(PageSize).Do(ctx.Req.Context())
		if err == nil {
			result := makeRepoResult(res, "")
			ctx.JSON(200, result)
		} else {
			log.Info("query es error," + err.Error())
			ctx.JSON(200, "")
		}
	}

}

func makeDatasetResult(sRes *elastic.SearchResult, Key string) *SearchRes {
	total := sRes.Hits.TotalHits.Value
	result := make([]map[string]interface{}, 0)

	for i, hit := range sRes.Hits.Hits {
		log.Info("this is dataset query " + fmt.Sprint(i) + " result.")
		recordSource := make(map[string]interface{})
		source, err := hit.Source.MarshalJSON()

		if err == nil {
			err = json.Unmarshal(source, &recordSource)
			if err == nil {
				record := make(map[string]interface{})
				record["id"] = recordSource["id"]
				userId := recordSource["user_id"].(int64)
				user, errUser := models.GetUserByID(userId)
				if errUser == nil {
					record["owerName"] = user.GetDisplayName()
					record["avatar"] = user.RelAvatarLink()
				}
				record["title"] = recordSource["title"]
				record["category"] = recordSource["category"]
				desc := recordSource["description"].(string)
				record["description"] = dealLongText(desc, Key, hit.MatchedQueries)
				record["download_times"] = recordSource["download_times"]
				record["created_unix"] = recordSource["created_unix"]
				result = append(result, record)
			} else {
				log.Info("deal dataset source error," + err.Error())
			}
		} else {
			log.Info("deal dataset source error," + err.Error())
		}
	}

	returnObj := &SearchRes{
		Total:  total,
		Result: result,
	}

	return returnObj
}

func searchIssue(ctx *context.Context, TableName string, Key string, Page int, PageSize int) {

	/*
	   任务，合并请求  ES名称：issue-es-index
	   搜索:
	   name character varying(255) ,  标题
	   content text,  内容
	   comment text,  评论
	   排序:
	   updated_unix
	*/
	SortBy := ctx.Query("SortBy")
	if SortBy == "" {
		SortBy = "updated_unix.keyword"
	}
	ascending := ctx.QueryBool("Ascending")
	boolQ := elastic.NewBoolQuery()
	if Key != "" {
		nameQuery := elastic.NewMatchQuery("name", Key).Boost(2).QueryName("f_first")
		contentQuery := elastic.NewMatchQuery("content", Key).Boost(1.5).QueryName("f_second")
		commentQuery := elastic.NewMatchQuery("comment", Key).Boost(1).QueryName("f_third")
		boolQ.Should(nameQuery, contentQuery, commentQuery)
	}
	isIssueQuery := elastic.NewTermQuery("is_pull", "f")
	boolQ.Must(isIssueQuery)
	res, err := client.Search(TableName).Query(boolQ).Sort(SortBy, ascending).From((Page - 1) * PageSize).Size(PageSize).Do(ctx.Req.Context())
	if err == nil {
		result := makeIssueResult(res, Key)
		ctx.JSON(200, result)
	} else {
		log.Info("query es error," + err.Error())
	}

}

func makeIssueResult(sRes *elastic.SearchResult, Key string) *SearchRes {
	total := sRes.Hits.TotalHits.Value
	result := make([]map[string]interface{}, 0)

	for i, hit := range sRes.Hits.Hits {
		log.Info("this is issue query " + fmt.Sprint(i) + " result.")
		recordSource := make(map[string]interface{})
		source, err := hit.Source.MarshalJSON()

		if err == nil {
			err = json.Unmarshal(source, &recordSource)
			if err == nil {
				record := make(map[string]interface{})
				record["id"] = recordSource["id"]
				record["repo_id"] = recordSource["repo_id"]
				log.Info("recordSource[\"repo_id\"]=" + fmt.Sprint(recordSource["repo_id"]))
				repo, errRepo := models.GetRepositoryByID(recordSource["repo_id"].(int64))
				if errRepo == nil {
					record["repoUrl"] = repo.FullName()
					record["avatar"] = repo.RelAvatarLink()
				}
				record["name"] = recordSource["name"]
				desc := recordSource["content"].(string)
				record["content"] = dealLongText(desc, Key, hit.MatchedQueries)
				if recordSource["pr_id"] != nil {
					record["pr_id"] = recordSource["pr_id"]
				}

				desc = recordSource["comment"].(string)
				record["comment"] = dealLongText(desc, Key, hit.MatchedQueries)
				record["num_comments"] = recordSource["num_comments"]
				record["updated_unix"] = recordSource["updated_unix"]
				result = append(result, record)
			} else {
				log.Info("deal issue source error," + err.Error())
			}
		} else {
			log.Info("deal issue source error," + err.Error())
		}
	}

	returnObj := &SearchRes{
		Total:  total,
		Result: result,
	}

	return returnObj
}

func searchPR(ctx *context.Context, TableName string, Key string, Page int, PageSize int) {

	/*
	   任务，合并请求  ES名称：issue-es-index
	   搜索:
	   name character varying(255) ,  标题
	   content text,  内容
	   comment text,  评论
	   排序:
	   updated_unix
	*/
	SortBy := ctx.Query("SortBy")
	if SortBy == "" {
		SortBy = "updated_unix.keyword"
	}
	ascending := ctx.QueryBool("Ascending")
	boolQ := elastic.NewBoolQuery()
	if Key != "" {
		nameQuery := elastic.NewMatchQuery("name", Key).Boost(1024).QueryName("f_first")
		contentQuery := elastic.NewMatchQuery("content", Key).Boost(1.5).QueryName("f_second")
		commentQuery := elastic.NewMatchQuery("comment", Key).Boost(1).QueryName("f_third")
		boolQ.Should(nameQuery, contentQuery, commentQuery)
	}
	isIssueQuery := elastic.NewTermQuery("is_pull", "t")
	boolQ.Must(isIssueQuery)
	res, err := client.Search(TableName).Query(boolQ).Sort(SortBy, ascending).From((Page - 1) * PageSize).Size(PageSize).Do(ctx.Req.Context())
	if err == nil {
		result := makeIssueResult(res, Key)
		ctx.JSON(200, result)
	} else {
		log.Info("query es error," + err.Error())
	}

}