1// Copyright 2012-present Oliver Eilhard. All rights reserved.
2// Use of this source code is governed by a MIT-license.
3// See http://olivere.mit-license.org/license.txt for details.
4
5package elastic
6
7import "errors"
8
9// CompletionSuggester is a fast suggester for e.g. type-ahead completion.
10//
11// See https://www.elastic.co/guide/en/elasticsearch/reference/6.7/search-suggesters-completion.html
12// for more details.
13type CompletionSuggester struct {
14	Suggester
15	name           string
16	text           string
17	prefix         string
18	regex          string
19	field          string
20	analyzer       string
21	size           *int
22	shardSize      *int
23	contextQueries []SuggesterContextQuery
24	payload        interface{}
25
26	fuzzyOptions   *FuzzyCompletionSuggesterOptions
27	regexOptions   *RegexCompletionSuggesterOptions
28	skipDuplicates *bool
29}
30
31// Creates a new completion suggester.
32func NewCompletionSuggester(name string) *CompletionSuggester {
33	return &CompletionSuggester{
34		name: name,
35	}
36}
37
38func (q *CompletionSuggester) Name() string {
39	return q.name
40}
41
42func (q *CompletionSuggester) Text(text string) *CompletionSuggester {
43	q.text = text
44	return q
45}
46
47func (q *CompletionSuggester) Prefix(prefix string) *CompletionSuggester {
48	q.prefix = prefix
49	return q
50}
51
52func (q *CompletionSuggester) PrefixWithEditDistance(prefix string, editDistance interface{}) *CompletionSuggester {
53	q.prefix = prefix
54	q.fuzzyOptions = NewFuzzyCompletionSuggesterOptions().EditDistance(editDistance)
55	return q
56}
57
58func (q *CompletionSuggester) PrefixWithOptions(prefix string, options *FuzzyCompletionSuggesterOptions) *CompletionSuggester {
59	q.prefix = prefix
60	q.fuzzyOptions = options
61	return q
62}
63
64func (q *CompletionSuggester) FuzzyOptions(options *FuzzyCompletionSuggesterOptions) *CompletionSuggester {
65	q.fuzzyOptions = options
66	return q
67}
68
69func (q *CompletionSuggester) Fuzziness(fuzziness interface{}) *CompletionSuggester {
70	if q.fuzzyOptions == nil {
71		q.fuzzyOptions = NewFuzzyCompletionSuggesterOptions()
72	}
73	q.fuzzyOptions = q.fuzzyOptions.EditDistance(fuzziness)
74	return q
75}
76
77func (q *CompletionSuggester) Regex(regex string) *CompletionSuggester {
78	q.regex = regex
79	return q
80}
81
82func (q *CompletionSuggester) RegexWithOptions(regex string, options *RegexCompletionSuggesterOptions) *CompletionSuggester {
83	q.regex = regex
84	q.regexOptions = options
85	return q
86}
87
88func (q *CompletionSuggester) RegexOptions(options *RegexCompletionSuggesterOptions) *CompletionSuggester {
89	q.regexOptions = options
90	return q
91}
92
93func (q *CompletionSuggester) SkipDuplicates(skipDuplicates bool) *CompletionSuggester {
94	q.skipDuplicates = &skipDuplicates
95	return q
96}
97
98func (q *CompletionSuggester) Field(field string) *CompletionSuggester {
99	q.field = field
100	return q
101}
102
103func (q *CompletionSuggester) Analyzer(analyzer string) *CompletionSuggester {
104	q.analyzer = analyzer
105	return q
106}
107
108func (q *CompletionSuggester) Size(size int) *CompletionSuggester {
109	q.size = &size
110	return q
111}
112
113func (q *CompletionSuggester) ShardSize(shardSize int) *CompletionSuggester {
114	q.shardSize = &shardSize
115	return q
116}
117
118func (q *CompletionSuggester) ContextQuery(query SuggesterContextQuery) *CompletionSuggester {
119	q.contextQueries = append(q.contextQueries, query)
120	return q
121}
122
123func (q *CompletionSuggester) ContextQueries(queries ...SuggesterContextQuery) *CompletionSuggester {
124	q.contextQueries = append(q.contextQueries, queries...)
125	return q
126}
127
128// completionSuggesterRequest is necessary because the order in which
129// the JSON elements are routed to Elasticsearch is relevant.
130// We got into trouble when using plain maps because the text element
131// needs to go before the completion element.
132type completionSuggesterRequest struct {
133	Text       string      `json:"text,omitempty"`
134	Prefix     string      `json:"prefix,omitempty"`
135	Regex      string      `json:"regex,omitempty"`
136	Completion interface{} `json:"completion,omitempty"`
137}
138
139// Source creates the JSON data for the completion suggester.
140func (q *CompletionSuggester) Source(includeName bool) (interface{}, error) {
141	cs := &completionSuggesterRequest{}
142
143	if q.text != "" {
144		cs.Text = q.text
145	}
146	if q.prefix != "" {
147		cs.Prefix = q.prefix
148	}
149	if q.regex != "" {
150		cs.Regex = q.regex
151	}
152
153	suggester := make(map[string]interface{})
154	cs.Completion = suggester
155
156	if q.analyzer != "" {
157		suggester["analyzer"] = q.analyzer
158	}
159	if q.field != "" {
160		suggester["field"] = q.field
161	}
162	if q.size != nil {
163		suggester["size"] = *q.size
164	}
165	if q.shardSize != nil {
166		suggester["shard_size"] = *q.shardSize
167	}
168	switch len(q.contextQueries) {
169	case 0:
170	case 1:
171		src, err := q.contextQueries[0].Source()
172		if err != nil {
173			return nil, err
174		}
175		suggester["contexts"] = src
176	default:
177		ctxq := make(map[string]interface{})
178		for _, query := range q.contextQueries {
179			src, err := query.Source()
180			if err != nil {
181				return nil, err
182			}
183			// Merge the dictionary into ctxq
184			m, ok := src.(map[string]interface{})
185			if !ok {
186				return nil, errors.New("elastic: context query is not a map")
187			}
188			for k, v := range m {
189				ctxq[k] = v
190			}
191		}
192		suggester["contexts"] = ctxq
193	}
194
195	// Fuzzy options
196	if q.fuzzyOptions != nil {
197		src, err := q.fuzzyOptions.Source()
198		if err != nil {
199			return nil, err
200		}
201		suggester["fuzzy"] = src
202	}
203
204	// Regex options
205	if q.regexOptions != nil {
206		src, err := q.regexOptions.Source()
207		if err != nil {
208			return nil, err
209		}
210		suggester["regex"] = src
211	}
212
213	if q.skipDuplicates != nil {
214		suggester["skip_duplicates"] = *q.skipDuplicates
215	}
216
217	// TODO(oe) Add completion-suggester specific parameters here
218
219	if !includeName {
220		return cs, nil
221	}
222
223	source := make(map[string]interface{})
224	source[q.name] = cs
225	return source, nil
226}
227
228// -- Fuzzy options --
229
230// FuzzyCompletionSuggesterOptions represents the options for fuzzy completion suggester.
231type FuzzyCompletionSuggesterOptions struct {
232	editDistance          interface{}
233	transpositions        *bool
234	minLength             *int
235	prefixLength          *int
236	unicodeAware          *bool
237	maxDeterminizedStates *int
238}
239
240// NewFuzzyCompletionSuggesterOptions initializes a new FuzzyCompletionSuggesterOptions instance.
241func NewFuzzyCompletionSuggesterOptions() *FuzzyCompletionSuggesterOptions {
242	return &FuzzyCompletionSuggesterOptions{}
243}
244
245// EditDistance specifies the maximum number of edits, e.g. a number like "1" or "2"
246// or a string like "0..2" or ">5".
247//
248// See https://www.elastic.co/guide/en/elasticsearch/reference/6.7/common-options.html#fuzziness
249// for details.
250func (o *FuzzyCompletionSuggesterOptions) EditDistance(editDistance interface{}) *FuzzyCompletionSuggesterOptions {
251	o.editDistance = editDistance
252	return o
253}
254
255// Transpositions, if set to true, are counted as one change instead of two (defaults to true).
256func (o *FuzzyCompletionSuggesterOptions) Transpositions(transpositions bool) *FuzzyCompletionSuggesterOptions {
257	o.transpositions = &transpositions
258	return o
259}
260
261// MinLength represents the minimum length of the input before fuzzy suggestions are returned (defaults to 3).
262func (o *FuzzyCompletionSuggesterOptions) MinLength(minLength int) *FuzzyCompletionSuggesterOptions {
263	o.minLength = &minLength
264	return o
265}
266
267// PrefixLength represents the minimum length of the input, which is not checked for
268// fuzzy alternatives (defaults to 1).
269func (o *FuzzyCompletionSuggesterOptions) PrefixLength(prefixLength int) *FuzzyCompletionSuggesterOptions {
270	o.prefixLength = &prefixLength
271	return o
272}
273
274// UnicodeAware, if true, all measurements (like fuzzy edit distance, transpositions, and lengths)
275// are measured in Unicode code points instead of in bytes. This is slightly slower than
276// raw bytes, so it is set to false by default.
277func (o *FuzzyCompletionSuggesterOptions) UnicodeAware(unicodeAware bool) *FuzzyCompletionSuggesterOptions {
278	o.unicodeAware = &unicodeAware
279	return o
280}
281
282// MaxDeterminizedStates is currently undocumented in Elasticsearch. It represents
283// the maximum automaton states allowed for fuzzy expansion.
284func (o *FuzzyCompletionSuggesterOptions) MaxDeterminizedStates(max int) *FuzzyCompletionSuggesterOptions {
285	o.maxDeterminizedStates = &max
286	return o
287}
288
289// Source creates the JSON data.
290func (o *FuzzyCompletionSuggesterOptions) Source() (interface{}, error) {
291	out := make(map[string]interface{})
292
293	if o.editDistance != nil {
294		out["fuzziness"] = o.editDistance
295	}
296	if o.transpositions != nil {
297		out["transpositions"] = *o.transpositions
298	}
299	if o.minLength != nil {
300		out["min_length"] = *o.minLength
301	}
302	if o.prefixLength != nil {
303		out["prefix_length"] = *o.prefixLength
304	}
305	if o.unicodeAware != nil {
306		out["unicode_aware"] = *o.unicodeAware
307	}
308	if o.maxDeterminizedStates != nil {
309		out["max_determinized_states"] = *o.maxDeterminizedStates
310	}
311
312	return out, nil
313}
314
315// -- Regex options --
316
317// RegexCompletionSuggesterOptions represents the options for regex completion suggester.
318type RegexCompletionSuggesterOptions struct {
319	flags                 interface{} // string or int
320	maxDeterminizedStates *int
321}
322
323// NewRegexCompletionSuggesterOptions initializes a new RegexCompletionSuggesterOptions instance.
324func NewRegexCompletionSuggesterOptions() *RegexCompletionSuggesterOptions {
325	return &RegexCompletionSuggesterOptions{}
326}
327
328// Flags represents internal regex flags.
329// Possible flags are ALL (default), ANYSTRING, COMPLEMENT, EMPTY, INTERSECTION, INTERVAL, or NONE.
330//
331// See https://www.elastic.co/guide/en/elasticsearch/reference/6.7/search-suggesters-completion.html#regex
332// for details.
333func (o *RegexCompletionSuggesterOptions) Flags(flags interface{}) *RegexCompletionSuggesterOptions {
334	o.flags = flags
335	return o
336}
337
338// MaxDeterminizedStates represents the maximum automaton states allowed for regex expansion.
339//
340// See https://www.elastic.co/guide/en/elasticsearch/reference/6.7/search-suggesters-completion.html#regex
341// for details.
342func (o *RegexCompletionSuggesterOptions) MaxDeterminizedStates(max int) *RegexCompletionSuggesterOptions {
343	o.maxDeterminizedStates = &max
344	return o
345}
346
347// Source creates the JSON data.
348func (o *RegexCompletionSuggesterOptions) Source() (interface{}, error) {
349	out := make(map[string]interface{})
350
351	if o.flags != nil {
352		out["flags"] = o.flags
353	}
354	if o.maxDeterminizedStates != nil {
355		out["max_determinized_states"] = *o.maxDeterminizedStates
356	}
357
358	return out, nil
359}
360