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 {
		log.Info("es init error.")
		//panic(err)
	}
}

func Search(ctx *context.Context) {
	keyword := strings.Trim(ctx.Query("q"), " ")
	ctx.Data["Keyword"] = keyword
	ctx.Data["SortType"] = "newest"

	ctx.HTML(200, "explore/search_new")
}

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

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

func searchRepoByLabel(ctx *context.Context, 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("repository-es-index").Query(boolQ).SortBy(elastic.NewScoreSort(), elastic.NewFieldSort(SortBy).Order(ascending)).From((Page - 1) * PageSize).Size(PageSize).Highlight(queryHighlight("topics")).Do(ctx.Req.Context())
		if err == nil {
			searchJson, _ := json.Marshal(res)
			log.Info("searchJson=" + string(searchJson))
			result := makeRepoResult(res, "", false)
			ctx.JSON(200, result)
		} else {
			log.Info("query es error," + err.Error())
			ctx.JSON(200, "")
		}
	}
	ctx.JSON(200, "")
}

func searchRepo(ctx *context.Context, TableName string, Key string, Page int, PageSize int, OnlyReturnNum bool) {
	/*
			   项目， 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).Highlight(queryHighlight("name", "description", "topics")).Do(ctx.Req.Context())
		if err == nil {
			searchJson, _ := json.Marshal(res)
			log.Info("searchJson=" + string(searchJson))
			result := makeRepoResult(res, Key, OnlyReturnNum)
			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 {
			searchJson, _ := json.Marshal(res)
			log.Info("searchJson=" + string(searchJson))
			result := makeRepoResult(res, "", OnlyReturnNum)
			ctx.JSON(200, result)
		} else {
			log.Info("query es error," + err.Error())
			ctx.JSON(200, "")
		}
	}
}

func makeRepoResult(sRes *elastic.SearchResult, Key string, OnlyReturnNum bool) *SearchRes {
	total := sRes.Hits.TotalHits.Value
	result := make([]map[string]interface{}, 0)
	if !OnlyReturnNum {
		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["id"] = hit.Id
					record["name"] = getLabelValue("name", recordSource, hit.Highlight)
					record["real_name"] = recordSource["name"]
					record["owner_name"] = recordSource["owner_name"]
					if recordSource["description"] != nil {
						desc := getLabelValue("description", recordSource, hit.Highlight)
						record["description"] = dealLongText(desc, Key, hit.MatchedQueries)
					} else {
						record["description"] = ""
					}
					if Key == "" {
						record["hightTopics"] = getLabelValue("topics", recordSource, hit.Highlight)
					}

					record["num_watches"] = recordSource["num_watches"]
					record["num_stars"] = recordSource["num_stars"]
					record["num_forks"] = recordSource["num_forks"]

					if recordSource["topics"] != nil {
						topicsStr := recordSource["topics"].(string)
						log.Info("topicsStr=" + topicsStr)
						if topicsStr != "null" {
							topicsStr = strings.Replace(topicsStr, "\"", "", -1)
							topicsStr = topicsStr[1 : len(topicsStr)-1]
							log.Info("record[\"topics\"]=" + topicsStr)
							topicstmp := strings.Split(topicsStr, ",")
							record["topics"] = topicstmp
						}
					}
					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, OnlyReturnNum bool) {
	/*
	   用户或者组织 ES名称： user-es-index
	   搜索:
	   name  ,  名称
	   full_name   全名
	   description  描述或者简介
	   排序:
	   created_unix
	   名称字母序
	*/
	SortBy := ctx.Query("SortBy")
	if SortBy == "" {
		SortBy = "created_unix.keyword"
	}
	ascending := ctx.QueryBool("Ascending")
	boolQ := elastic.NewBoolQuery()

	typeValue := 1
	if IsQueryUser {
		typeValue = 0
	}
	UserOrOrgQuery := elastic.NewTermQuery("type", typeValue)
	if Key != "" {
		boolKeyQ := elastic.NewBoolQuery()
		log.Info("issue Key=" + 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")
		boolKeyQ.Should(nameQuery, full_nameQuery, descriptionQuery)
		boolQ.Must(UserOrOrgQuery, boolKeyQ)
	} else {
		boolQ.Must(UserOrOrgQuery)
	}

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

func getLabelValue(key string, recordSource map[string]interface{}, searchHighliht elastic.SearchHitHighlight) string {
	if value, ok := searchHighliht[key]; !ok {
		if recordSource[key] != nil {
			return recordSource[key].(string)
		} else {
			return ""
		}
	} else {
		return value[0]
	}
}

func makeUserOrOrgResult(sRes *elastic.SearchResult, Key string, ctx *context.Context, OnlyReturnNum bool) *SearchRes {
	total := sRes.Hits.TotalHits.Value
	result := make([]map[string]interface{}, 0)
	if !OnlyReturnNum {
		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["id"] = hit.Id
					record["name"] = getLabelValue("name", recordSource, hit.Highlight)
					record["real_name"] = recordSource["name"]
					record["full_name"] = getLabelValue("full_name", recordSource, hit.Highlight)
					if recordSource["description"] != nil {
						desc := getLabelValue("description", recordSource, hit.Highlight)
						record["description"] = dealLongText(desc, Key, hit.MatchedQueries)
					} else {
						record["description"] = ""
					}
					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, OnlyReturnNum bool) {
	/*
	   数据集，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).Highlight(queryHighlight("title", "description", "file_name", "category")).Do(ctx.Req.Context())
		if err == nil {
			searchJson, _ := json.Marshal(res)
			log.Info("searchJson=" + string(searchJson))
			result := makeDatasetResult(res, Key, OnlyReturnNum)
			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 {
			searchJson, _ := json.Marshal(res)
			log.Info("searchJson=" + string(searchJson))
			result := makeDatasetResult(res, "", OnlyReturnNum)
			ctx.JSON(200, result)
		} else {
			log.Info("query es error," + err.Error())
			ctx.JSON(200, "")
		}
	}

}

func makeDatasetResult(sRes *elastic.SearchResult, Key string, OnlyReturnNum bool) *SearchRes {
	total := sRes.Hits.TotalHits.Value
	result := make([]map[string]interface{}, 0)
	if !OnlyReturnNum {
		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"] = hit.Id
					userIdStr := recordSource["user_id"].(string)
					userId, cerr := strconv.ParseInt(userIdStr, 10, 64)
					if cerr == nil {
						user, errUser := models.GetUserByID(userId)
						if errUser == nil {
							record["owerName"] = user.GetDisplayName()
							record["avatar"] = user.RelAvatarLink()
						}
					}
					setRepoInfo(recordSource, record)
					record["title"] = getLabelValue("title", recordSource, hit.Highlight)
					record["category"] = getLabelValue("category", recordSource, hit.Highlight)
					if recordSource["description"] != nil {
						desc := getLabelValue("description", recordSource, hit.Highlight)
						record["description"] = dealLongText(desc, Key, hit.MatchedQueries)
					} else {
						record["description"] = ""
					}
					record["file_name"] = recordSource["file_name"]
					record["task"] = recordSource["task"]
					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, OnlyReturnNum bool) {

	/*
	   任务，合并请求  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()
	isIssueQuery := elastic.NewTermQuery("is_pull", "f")

	if Key != "" {
		boolKeyQ := elastic.NewBoolQuery()
		log.Info("issue Key=" + 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")
		boolKeyQ.Should(nameQuery, contentQuery, commentQuery)
		boolQ.Must(isIssueQuery, boolKeyQ)
	} else {
		boolQ.Must(isIssueQuery)
	}

	res, err := client.Search(TableName).Query(boolQ).Sort(SortBy, ascending).From((Page - 1) * PageSize).Size(PageSize).Highlight(queryHighlight("name", "content", "comment")).Do(ctx.Req.Context())
	if err == nil {
		searchJson, _ := json.Marshal(res)
		log.Info("searchJson=" + string(searchJson))
		result := makeIssueResult(res, Key, OnlyReturnNum)
		ctx.JSON(200, result)
	} else {
		log.Info("query es error," + err.Error())
	}
}

func queryHighlight(names ...string) *elastic.Highlight {
	re := elastic.NewHighlight()
	for i := 0; i < len(names); i++ {
		field := &elastic.HighlighterField{
			Name: names[i],
		}
		re.Fields(field)
	}
	re.PreTags("<font color='red'>")
	re.PostTags("</font>")
	return re
}

func setRepoInfo(recordSource map[string]interface{}, record map[string]interface{}) {
	repoIdstr := recordSource["repo_id"].(string)
	repoId, cerr := strconv.ParseInt(repoIdstr, 10, 64)
	if cerr == nil {
		repo, errRepo := models.GetRepositoryByID(repoId)
		if errRepo == nil {
			log.Info("repo_url=" + repo.FullName())
			record["repoUrl"] = repo.FullName()
			record["avatar"] = repo.RelAvatarLink()
		} else {
			log.Info("repo err=" + errRepo.Error())
		}
	} else {
		log.Info("parse int err=" + cerr.Error())
	}
}

func makeIssueResult(sRes *elastic.SearchResult, Key string, OnlyReturnNum bool) *SearchRes {
	total := sRes.Hits.TotalHits.Value
	result := make([]map[string]interface{}, 0)
	if !OnlyReturnNum {
		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"] = hit.Id
					record["repo_id"] = recordSource["repo_id"]
					log.Info("recordSource[\"repo_id\"]=" + fmt.Sprint(recordSource["repo_id"]))
					setRepoInfo(recordSource, record)
					record["name"] = getLabelValue("name", recordSource, hit.Highlight)
					if recordSource["content"] != nil {
						desc := getLabelValue("content", recordSource, hit.Highlight)
						record["content"] = dealLongText(desc, Key, hit.MatchedQueries)
						if _, ok := hit.Highlight["content"]; !ok {
							if _, ok_comment := hit.Highlight["comment"]; ok_comment {
								desc := getLabelValue("comment", recordSource, hit.Highlight)
								record["content"] = dealLongText(desc, Key, hit.MatchedQueries)
							}
						}
					} else {
						if recordSource["comment"] != nil {
							desc := getLabelValue("comment", recordSource, hit.Highlight)
							record["content"] = dealLongText(desc, Key, hit.MatchedQueries)
						}
					}
					if recordSource["pr_id"] != nil {
						record["pr_id"] = recordSource["pr_id"]
					}

					record["num_comments"] = recordSource["num_comments"]
					record["is_closed"] = recordSource["is_closed"]
					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, OnlyReturnNum bool) {

	/*
	   任务，合并请求  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()
	isPRQuery := elastic.NewTermQuery("is_pull", "t")

	if Key != "" {
		boolKeyQ := elastic.NewBoolQuery()
		log.Info("issue Key=" + 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")
		boolKeyQ.Should(nameQuery, contentQuery, commentQuery)
		boolQ.Must(isPRQuery, boolKeyQ)
	} else {
		boolQ.Must(isPRQuery)
	}
	res, err := client.Search(TableName).Query(boolQ).Sort(SortBy, ascending).From((Page - 1) * PageSize).Size(PageSize).Highlight(queryHighlight("name", "content", "comment")).Do(ctx.Req.Context())
	if err == nil {
		result := makeIssueResult(res, Key, OnlyReturnNum)
		ctx.JSON(200, result)
	} else {
		log.Info("query es error," + err.Error())
	}

}