1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71
72. 72
73. 73
74. 74
75. 75
76. 76
77. 77
78. 78
79. 79
80. 80
81. 81
82. 82
83. 83
84. 84
85. 85
86. 86
87. 87
88. 88
89. 89
90. 90
91. 91
92. 92
93. 93
94. 94
95. 95
96. 96
97. 97
98. 98
99. 99
100. 100
101. 101
102. 102
103. 103
104. 104
105. 105
106. 106
107. 107
108. 108
109. 109
110. 110
111. 111
112. 112
113. 113
114. 114
115. 115
116. 116
117. 117
118. 118
119. 119
120. 120
121. 121
122. 122
123. 123
124. 124

// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

import "errors"

// SuggesterContextQuery is used to define context information within
// a suggestion request.
type SuggesterContextQuery interface {
	Source() (interface{}, error)
}

// ContextSuggester is a fast suggester for e.g. type-ahead completion that supports filtering and boosting based on contexts.
// See https://www.elastic.co/guide/en/elasticsearch/reference/7.0/suggester-context.html
// for more details.
type ContextSuggester struct {
	Suggester
	name           string
	prefix         string
	field          string
	size           *int
	contextQueries []SuggesterContextQuery
}

// Creates a new context suggester.
func NewContextSuggester(name string) *ContextSuggester {
	return &ContextSuggester{
		name:           name,
		contextQueries: make([]SuggesterContextQuery, 0),
	}
}

func (q *ContextSuggester) Name() string {
	return q.name
}

func (q *ContextSuggester) Prefix(prefix string) *ContextSuggester {
	q.prefix = prefix
	return q
}

func (q *ContextSuggester) Field(field string) *ContextSuggester {
	q.field = field
	return q
}

func (q *ContextSuggester) Size(size int) *ContextSuggester {
	q.size = &size
	return q
}

func (q *ContextSuggester) ContextQuery(query SuggesterContextQuery) *ContextSuggester {
	q.contextQueries = append(q.contextQueries, query)
	return q
}

func (q *ContextSuggester) ContextQueries(queries ...SuggesterContextQuery) *ContextSuggester {
	q.contextQueries = append(q.contextQueries, queries...)
	return q
}

// contextSuggesterRequest is necessary because the order in which
// the JSON elements are routed to Elasticsearch is relevant.
// We got into trouble when using plain maps because the text element
// needs to go before the completion element.
type contextSuggesterRequest struct {
	Prefix     string      `json:"prefix"`
	Completion interface{} `json:"completion"`
}

// Creates the source for the context suggester.
func (q *ContextSuggester) Source(includeName bool) (interface{}, error) {
	cs := &contextSuggesterRequest{}

	if q.prefix != "" {
		cs.Prefix = q.prefix
	}

	suggester := make(map[string]interface{})
	cs.Completion = suggester

	if q.field != "" {
		suggester["field"] = q.field
	}
	if q.size != nil {
		suggester["size"] = *q.size
	}
	switch len(q.contextQueries) {
	case 0:
	case 1:
		src, err := q.contextQueries[0].Source()
		if err != nil {
			return nil, err
		}
		suggester["contexts"] = src
	default:
		ctxq := make(map[string]interface{})
		for _, query := range q.contextQueries {
			src, err := query.Source()
			if err != nil {
				return nil, err
			}
			// Merge the dictionary into ctxq
			m, ok := src.(map[string]interface{})
			if !ok {
				return nil, errors.New("elastic: context query is not a map")
			}
			for k, v := range m {
				ctxq[k] = v
			}
		}
		suggester["contexts"] = ctxq
	}

	if !includeName {
		return cs, nil
	}

	source := make(map[string]interface{})
	source[q.name] = cs
	return source, nil
}