modelRT/model/redis_recommend.go

// Package model define model struct of model runtime service
package model

import (
	"context"
	"fmt"
	"math"
	"strings"

	"modelRT/constants"
	"modelRT/diagram"
	"modelRT/logger"

	"github.com/RediSearch/redisearch-go/v2/redisearch"
	redigo "github.com/gomodule/redigo/redis"
)

var ac *redisearch.Autocompleter

// InitAutocompleterWithPool define func of initialize the Autocompleter with redigo pool
func InitAutocompleterWithPool(pool *redigo.Pool) {
	ac = redisearch.NewAutocompleterFromPool(pool, constants.RedisSearchDictName)
}

// RedisSearchRecommend define func of redis search by input string and return recommend results
func RedisSearchRecommend(ctx context.Context, input string) ([]string, bool, error) {
	rdb := diagram.GetRedisClientInstance()

	if input == "" {
		// 返回所有 grid 名
		return getAllGridKeys(ctx, constants.RedisAllGridSetKey)
	}

	inputSlice := strings.Split(input, ".")
	inputSliceLen := len(inputSlice)
	originInputLen := len(inputSlice)

	switch inputSliceLen {
	case 1:
		// TODO 优化成NewSet的形式
		gridExist, err := rdb.SIsMember(ctx, constants.RedisAllGridSetKey, input).Result()
		if err != nil {
			logger.Error(ctx, "check grid key exist failed ", "grid_key", input, "error", err)
			return []string{}, false, err
		}

		searchInput := input
		inputLen := inputSliceLen
		for inputLen != 0 && !gridExist {
			results, err := ac.SuggestOpts(searchInput, redisearch.SuggestOptions{
				Num:          math.MaxInt16,
				Fuzzy:        true,
				WithScores:   false,
				WithPayloads: false,
			})
			if err != nil {
				logger.Error(ctx, "query info by fuzzy failed", "query_key", input, "error", err)
				return []string{}, false, err
			}

			if len(results) == 0 {
				// TODO 考虑使用其他方式代替for 循环退一字节的查询方式
				searchInput = searchInput[:len(searchInput)-1]
				inputLen = len(searchInput)
				continue
			}

			var recommends []string
			for _, result := range results {
				termSlice := strings.Split(result.Term, ".")
				if len(termSlice) <= originInputLen {
					recommends = append(recommends, result.Term)
				}
			}
			// 返回模糊查询结果
			return recommends, true, nil
		}

		// 处理 input 不为空、不含有.并且 input 是一个完整的 grid key 的情况
		if strings.HasSuffix(input, ".") == false {
			recommend := input + "."
			return []string{recommend}, false, nil
		}
	default:
		lastInput := inputSlice[inputSliceLen-1]
		// 判断 queryKey 是否是空值，空值则返回上一级别下的所有key
		if lastInput == "" {
			setKey := getCombinedConstantsKeyByLength(inputSlice[inputSliceLen-2], inputSliceLen)
			targetSet := diagram.NewRedisSet(ctx, setKey, 10, true)
			keys, err := targetSet.SMembers(setKey)
			if err != nil {
				logger.Error(ctx, "get all recommend key by setKey failed", "set_key", setKey, "error", err)
				return []string{}, false, fmt.Errorf("get all recommend key by setKey failed,%w", err)
			}

			var results []string
			for _, key := range keys {
				result := input + key
				results = append(results, result)
			}
			return results, false, nil
		}

		setKey := getCombinedConstantsKeyByLength(inputSlice[inputSliceLen-2], inputSliceLen)
		targetSet := diagram.NewRedisSet(ctx, setKey, 10, true)
		exist, err := targetSet.SIsMember(setKey, lastInput)
		if err != nil {
			logger.Error(ctx, "check keys exist failed", "set_key", setKey, "query_key", lastInput, "error", err)
			return []string{}, false, fmt.Errorf("check keys failed,%w", err)
		}

		searchInput := input
		inputLen := len(searchInput)
		for inputLen != 0 && !exist {
			logger.Info(ctx, "use fuzzy query", "input", input)
			results, err := ac.SuggestOpts(searchInput, redisearch.SuggestOptions{
				Num:          math.MaxInt16,
				Fuzzy:        true,
				WithScores:   false,
				WithPayloads: false,
			})
			if err != nil {
				logger.Error(ctx, "query info by fuzzy failed", "query_key", input, "error", err)
				return []string{}, false, err
			}
			if len(results) == 0 {
				searchInput = input[:inputLen-1]
				inputLen = len(searchInput)
				continue
			}

			var terms []string
			for _, result := range results {
				terms = append(terms, result.Term)
			}
			// 返回模糊查询结果
			return terms, true, nil
		}
		return []string{input}, false, nil
	}
	return []string{}, false, nil
}

func getAllGridKeys(ctx context.Context, setKey string) ([]string, bool, error) {
	// 从redis set 中获取所有的 grid key
	gridSets := diagram.NewRedisSet(ctx, setKey, 10, true)
	keys, err := gridSets.SMembers("grid_keys")
	if err != nil {
		return []string{}, false, fmt.Errorf("get all root keys failed, error: %v", err)
	}
	return keys, false, nil
}

func getSpecificZoneKeys(ctx context.Context, input string) ([]string, bool, error) {
	setKey := fmt.Sprintf(constants.RedisSpecGridZoneSetKey, input)
	zoneSets := diagram.NewRedisSet(ctx, setKey, 10, true)
	keys, err := zoneSets.SMembers(setKey)
	if err != nil {
		return []string{}, false, fmt.Errorf("get all root keys failed, error: %v", err)
	}
	var results []string
	for _, key := range keys {
		result := input + "." + key
		results = append(results, result)
	}
	return results, false, nil
}

func getConstantsKeyByLength(inputLen int) string {
	switch inputLen {
	case 1:
		return constants.RedisAllGridSetKey
	case 2:
		return constants.RedisAllZoneSetKey
	case 3:
		return constants.RedisAllStationSetKey
	case 4:
		return constants.RedisAllComponentSetKey
	default:
		return constants.RedisAllGridSetKey
	}
}

func getCombinedConstantsKeyByLength(key string, inputLen int) string {
	switch inputLen {
	case 2:
		return fmt.Sprintf(constants.RedisSpecGridZoneSetKey, key)
	case 3:
		return fmt.Sprintf(constants.RedisSpecZoneStationSetKey, key)
	case 4:
		return fmt.Sprintf(constants.RedisSpecStationComponentSetKey, key)
	default:
		return constants.RedisAllGridSetKey
	}
}

// GetLongestCommonPrefixLength define func of get longest common prefix length between two strings
func GetLongestCommonPrefixLength(input string, recommendResult string) int {
	if input == "" {
		return 0
	}

	minLen := min(len(input), len(recommendResult))

	for i := range minLen {
		if input[i] != recommendResult[i] {
			return i
		}
	}
	return minLen
}