Skip to content

Commit

Permalink
Move elastic utils from goflow
Browse files Browse the repository at this point in the history
  • Loading branch information
rowanseymour committed May 20, 2024
1 parent 69880ce commit fa6f40a
Show file tree
Hide file tree
Showing 4 changed files with 238 additions and 0 deletions.
132 changes: 132 additions & 0 deletions elastic/query.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
package elastic

type Query map[string]any

// Not is a shortcut for an ids query
func Ids(values ...string) Query {
return Query{"ids": map[string]any{"values": values}}
}

// Term is a shortcut for a term query
func Term(field string, value any) Query {
return Query{"term": map[string]any{field: value}}
}

// Exists is a shortcut for an exists query
func Exists(field string) Query {
return Query{"exists": map[string]any{"field": field}}
}

// Match is a shortcut for a match query
func Match(field string, value any) Query {
return Query{"match": map[string]any{field: map[string]any{"query": value}}}
}

// MatchPhrase is a shortcut for a match_phrase query
func MatchPhrase(field, value string) Query {
return Query{"match_phrase": map[string]any{field: map[string]any{"query": value}}}
}

// GreaterThan is a shortcut for a range query where x > value
func GreaterThan(field string, value any) Query {
return Query{
"range": map[string]any{
field: map[string]any{
"from": value,
"include_lower": false,
"include_upper": true,
"to": nil,
},
},
}
}

// GreaterThanOrEqual is a shortcut for a range query where x >= value
func GreaterThanOrEqual(field string, value any) Query {
return Query{
"range": map[string]any{
field: map[string]any{
"from": value,
"include_lower": true,
"include_upper": true,
"to": nil,
},
},
}
}

// LessThan is a shortcut for a range query where x < value
func LessThan(field string, value any) Query {
return Query{
"range": map[string]any{
field: map[string]any{
"from": nil,
"include_lower": true,
"include_upper": false,
"to": value,
},
},
}
}

// LessThanOrEqual is a shortcut for a range query where x <= value
func LessThanOrEqual(field string, value any) Query {
return Query{
"range": map[string]any{
field: map[string]any{
"from": nil,
"include_lower": true,
"include_upper": true,
"to": value,
},
},
}
}

// Between is a shortcut for a range query where from <= x < to
func Between(field string, from, to any) Query {
return Query{
"range": map[string]any{
field: map[string]any{
"from": from,
"include_lower": true,
"include_upper": false,
"to": to,
},
},
}
}

// Any is a shortcut for a bool query with a should clause
func Any(queries ...Query) Query {
return Query{"bool": map[string]any{"should": queries}}
}

// All is a shortcut for a bool query with a must clause
func All(queries ...Query) Query {
return Query{"bool": map[string]any{"must": queries}}
}

// Not is a shortcut for a bool query with a must_not clause
func Not(query Query) Query {
return Query{"bool": map[string]any{"must_not": query}}
}

// Bool is a shortcut for a bool query with multiple must and must_not clauses
func Bool(all []Query, none []Query) Query {
bq := map[string]any{}

if len(all) > 0 {
bq["must"] = all
}
if len(none) > 0 {
bq["must_not"] = none
}

return Query{"bool": bq}
}

// Nested is a shortcut for a nested query
func Nested(path string, query Query) Query {
return Query{"nested": map[string]any{"path": path, "query": query}}
}
56 changes: 56 additions & 0 deletions elastic/query_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
package elastic_test

import (
"testing"

"github.com/nyaruka/gocommon/elastic"
"github.com/nyaruka/gocommon/jsonx"
"github.com/stretchr/testify/assert"
)

func TestQuery(t *testing.T) {
tcs := []struct {
q elastic.Query
json []byte
}{
{elastic.Ids("235", "465", "787"), []byte(`{"ids": {"values": ["235", "465", "787"]}}`)},
{elastic.Term("age", 42), []byte(`{"term": {"age": 42}}`)},
{elastic.Exists("age"), []byte(`{"exists": {"field": "age"}}`)},
{elastic.Match("name", "Bob"), []byte(`{"match": {"name": {"query": "Bob"}}}`)},
{elastic.MatchPhrase("name", "Bob"), []byte(`{"match_phrase": {"name": {"query": "Bob"}}}`)},
{elastic.GreaterThan("age", 45), []byte(`{"range": {"age": {"from": 45, "include_lower": false, "include_upper": true, "to": null}}}`)},
{elastic.GreaterThanOrEqual("age", 45), []byte(`{"range": {"age": {"from": 45, "include_lower": true, "include_upper": true, "to": null}}}`)},
{elastic.LessThan("age", 45), []byte(`{"range": {"age": {"from": null, "include_lower": true, "include_upper": false, "to": 45}}}`)},
{elastic.LessThanOrEqual("age", 45), []byte(`{"range": {"age": {"from": null, "include_lower": true, "include_upper": true, "to": 45}}}`)},
{elastic.Between("age", 20, 45), []byte(`{"range": {"age": {"from": 20, "include_lower": true, "include_upper": false, "to": 45}}}`)},
{
elastic.Any(elastic.Ids("235"), elastic.Term("age", 42)),
[]byte(`{"bool": {"should": [{"ids": {"values": ["235"]}}, {"term": {"age": 42}}]}}`),
},
{
elastic.All(elastic.Ids("235"), elastic.Term("age", 42)),
[]byte(`{"bool": {"must": [{"ids": {"values": ["235"]}}, {"term": {"age": 42}}]}}`),
},
{
elastic.Not(elastic.Ids("235")),
[]byte(`{"bool": {"must_not": {"ids": {"values": ["235"]}}}}`),
},
{
elastic.Bool([]elastic.Query{elastic.Ids("235"), elastic.Term("age", 42)}, []elastic.Query{elastic.Exists("age")}),
[]byte(`{"bool": {"must": [{"ids": {"values": ["235"]}}, {"term": {"age": 42}}], "must_not": [{"exists": {"field": "age"}}]}}`),
},
{
elastic.Bool([]elastic.Query{}, []elastic.Query{elastic.Exists("age")}),
[]byte(`{"bool": {"must_not": [{"exists": {"field": "age"}}]}}`),
},
{
elastic.Bool([]elastic.Query{elastic.Ids("235"), elastic.Term("age", 42)}, []elastic.Query{}),
[]byte(`{"bool": {"must": [{"ids": {"values": ["235"]}}, {"term": {"age": 42}}]}}`),
},
{elastic.Nested("group", elastic.Term("group.id", 10)), []byte(`{"nested": {"path": "group", "query": {"term": {"group.id": 10}}}}`)},
}

for i, tc := range tcs {
assert.JSONEq(t, string(tc.json), string(jsonx.MustMarshal(tc.q)), "%d: elastic mismatch", i)
}
}
23 changes: 23 additions & 0 deletions elastic/sort.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
package elastic

type Sort map[string]any

// SortBy is a shortcut for a simple field sort
func SortBy(field string, ascending bool) Sort {
return Sort{field: map[string]any{"order": order(ascending)}}
}

// SortNested is a shortcut for a nested field sort
func SortNested(field string, filter Query, path string, ascending bool) Sort {
return Sort{field: map[string]any{
"nested": map[string]any{"filter": filter, "path": path},
"order": order(ascending),
}}
}

func order(asc bool) string {
if asc {
return "asc"
}
return "desc"
}
27 changes: 27 additions & 0 deletions elastic/sort_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
package elastic_test

import (
"testing"

"github.com/nyaruka/gocommon/elastic"
"github.com/nyaruka/gocommon/jsonx"
"github.com/stretchr/testify/assert"
)

func TestSort(t *testing.T) {
tcs := []struct {
q elastic.Sort
json []byte
}{
{elastic.SortBy("name", true), []byte(`{"name": {"order": "asc"}}`)},
{elastic.SortBy("name", false), []byte(`{"name": {"order": "desc"}}`)},
{
elastic.SortNested("age", elastic.Term("fields.field", "1234"), "fields", true),
[]byte(`{"age": {"nested": {"filter": {"term": {"fields.field": "1234"}}, "path": "fields"}, "order":"asc"}}`),
},
}

for i, tc := range tcs {
assert.JSONEq(t, string(tc.json), string(jsonx.MustMarshal(tc.q)), "%d: elastic mismatch", i)
}
}

0 comments on commit fa6f40a

Please sign in to comment.