Add Query and Sort map types

nyaruka · May 17, 2024 · 0c8234f · 0c8234f
1 parent 7114ff3
commit 0c8234f
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 49 deletions.
diff --git a/contactql/es/query.go b/contactql/es/query.go
@@ -26,15 +26,15 @@ var contactStatusCodes = map[string]string{
 }
 
 // ToElasticQuery converts a contactql query to an Elastic query
-func ToElasticQuery(env envs.Environment, mapper AssetMapper, query *contactql.ContactQuery) map[string]any {
+func ToElasticQuery(env envs.Environment, mapper AssetMapper, query *contactql.ContactQuery) elastic.Query {
 	if query.Resolver() == nil {
 		panic("can only convert queries parsed with a resolver")
 	}
 
 	return nodeToElastic(env, query.Resolver(), mapper, query.Root())
 }
 
-func nodeToElastic(env envs.Environment, resolver contactql.Resolver, mapper AssetMapper, node contactql.QueryNode) map[string]any {
+func nodeToElastic(env envs.Environment, resolver contactql.Resolver, mapper AssetMapper, node contactql.QueryNode) elastic.Query {
 	switch n := node.(type) {
 	case *contactql.BoolCombination:
 		return boolCombination(env, resolver, mapper, n)
@@ -45,8 +45,8 @@ func nodeToElastic(env envs.Environment, resolver contactql.Resolver, mapper Ass
 	}
 }
 
-func boolCombination(env envs.Environment, resolver contactql.Resolver, mapper AssetMapper, combination *contactql.BoolCombination) map[string]any {
-	queries := make([]map[string]any, len(combination.Children()))
+func boolCombination(env envs.Environment, resolver contactql.Resolver, mapper AssetMapper, combination *contactql.BoolCombination) elastic.Query {
+	queries := make([]elastic.Query, len(combination.Children()))
 	for i, child := range combination.Children() {
 		queries[i] = nodeToElastic(env, resolver, mapper, child)
 	}
@@ -58,7 +58,7 @@ func boolCombination(env envs.Environment, resolver contactql.Resolver, mapper A
 	return elastic.Any(queries...)
 }
 
-func condition(env envs.Environment, resolver contactql.Resolver, mapper AssetMapper, c *contactql.Condition) map[string]any {
+func condition(env envs.Environment, resolver contactql.Resolver, mapper AssetMapper, c *contactql.Condition) elastic.Query {
 	switch c.PropertyType() {
 	case contactql.PropertyTypeField:
 		return fieldCondition(env, resolver, c)
@@ -71,7 +71,7 @@ func condition(env envs.Environment, resolver contactql.Resolver, mapper AssetMa
 	}
 }
 
-func fieldCondition(env envs.Environment, resolver contactql.Resolver, c *contactql.Condition) map[string]any {
+func fieldCondition(env envs.Environment, resolver contactql.Resolver, c *contactql.Condition) elastic.Query {
 	field := resolver.ResolveField(c.PropertyKey())
 	fieldType := field.Type()
 	fieldQuery := elastic.Term("fields.field", field.UUID())
@@ -102,7 +102,7 @@ func fieldCondition(env envs.Environment, resolver contactql.Resolver, c *contac
 
 	} else if fieldType == assets.FieldTypeNumber {
 		value, _ := c.ValueAsNumber()
-		var query map[string]any
+		var query elastic.Query
 
 		switch c.Operator() {
 		case contactql.OpEqual:
@@ -130,7 +130,7 @@ func fieldCondition(env envs.Environment, resolver contactql.Resolver, c *contac
 	} else if fieldType == assets.FieldTypeDatetime {
 		value, _ := c.ValueAsDate(env)
 		start, end := dates.DayToUTCRange(value, value.Location())
-		var query map[string]any
+		var query elastic.Query
 
 		switch c.Operator() {
 		case contactql.OpEqual:
@@ -176,7 +176,7 @@ func fieldCondition(env envs.Environment, resolver contactql.Resolver, c *contac
 	panic(fmt.Sprintf("unsupported field type: %s", fieldType))
 }
 
-func attributeCondition(env envs.Environment, resolver contactql.Resolver, mapper AssetMapper, c *contactql.Condition) map[string]any {
+func attributeCondition(env envs.Environment, resolver contactql.Resolver, mapper AssetMapper, c *contactql.Condition) elastic.Query {
 	key := c.PropertyKey()
 	value := strings.ToLower(c.Value())
 
@@ -345,7 +345,7 @@ func attributeCondition(env envs.Environment, resolver contactql.Resolver, mappe
 	}
 }
 
-func schemeCondition(c *contactql.Condition) map[string]any {
+func schemeCondition(c *contactql.Condition) elastic.Query {
 	key := c.PropertyKey()
 	value := strings.ToLower(c.Value())
 
@@ -370,7 +370,7 @@ func schemeCondition(c *contactql.Condition) map[string]any {
 	}
 }
 
-func textAttributeQuery(c *contactql.Condition, name string, tx func(string) string) map[string]any {
+func textAttributeQuery(c *contactql.Condition, name string, tx func(string) string) elastic.Query {
 	value := tx(c.Value())
 
 	switch c.Operator() {
@@ -383,7 +383,7 @@ func textAttributeQuery(c *contactql.Condition, name string, tx func(string) str
 	}
 }
 
-func numericalAttributeQuery(c *contactql.Condition, name string) map[string]any {
+func numericalAttributeQuery(c *contactql.Condition, name string) elastic.Query {
 	value, _ := c.ValueAsNumber()
 
 	switch c.Operator() {

diff --git a/contactql/es/sort.go b/contactql/es/sort.go
@@ -11,10 +11,10 @@ import (
 )
 
 // ToElasticFieldSort returns the elastic FieldSort for the passed in sort by string
-func ToElasticFieldSort(sortBy string, resolver contactql.Resolver) (map[string]any, error) {
+func ToElasticFieldSort(sortBy string, resolver contactql.Resolver) (elastic.Sort, error) {
 	// default to most recent first by id
 	if sortBy == "" {
-		return elastic.Sort("id", false), nil
+		return elastic.SortBy("id", false), nil
 	}
 
 	// figure out if we are ascending or descending (default is ascending, can be changed with leading -)
@@ -29,12 +29,12 @@ func ToElasticFieldSort(sortBy string, resolver contactql.Resolver) (map[string]
 
 	// name needs to be sorted by keyword field
 	if property == contactql.AttributeName {
-		return elastic.Sort("name.keyword", ascending), nil
+		return elastic.SortBy("name.keyword", ascending), nil
 	}
 
 	// other attributes are straight sorts
 	if property == contactql.AttributeID || property == contactql.AttributeCreatedOn || property == contactql.AttributeLastSeenOn || property == contactql.AttributeLanguage {
-		return elastic.Sort(property, ascending), nil
+		return elastic.SortBy(property, ascending), nil
 	}
 
 	// we are sorting by a custom field

diff --git a/utils/elastic/query.go b/utils/elastic/query.go
@@ -1,53 +1,55 @@
 package elastic
 
+type Query map[string]any
+
 // Any is a shortcut for a bool query with a should clause
-func Any(queries ...map[string]any) map[string]any {
-	return map[string]any{"bool": map[string]any{"should": queries}}
+func Any(queries ...Query) Query {
+	return Query{"bool": Query{"should": queries}}
 }
 
 // All is a shortcut for a bool query with a must clause
-func All(queries ...map[string]any) map[string]any {
-	return map[string]any{"bool": map[string]any{"must": queries}}
+func All(queries ...Query) Query {
+	return Query{"bool": Query{"must": queries}}
 }
 
 // Not is a shortcut for a bool query with a must_not clause
-func Not(query map[string]any) map[string]any {
-	return map[string]any{"bool": map[string]any{"must_not": query}}
+func Not(query Query) Query {
+	return Query{"bool": Query{"must_not": query}}
 }
 
 // Not is a shortcut for an ids query
-func Ids(values ...string) map[string]any {
-	return map[string]any{"ids": map[string]any{"values": values}}
+func Ids(values ...string) Query {
+	return Query{"ids": Query{"values": values}}
 }
 
 // Term is a shortcut for a term query
-func Term(field string, value any) map[string]any {
-	return map[string]any{"term": map[string]any{field: value}}
+func Term(field string, value any) Query {
+	return Query{"term": Query{field: value}}
 }
 
 // Exists is a shortcut for an exists query
-func Exists(field string) map[string]any {
-	return map[string]any{"exists": map[string]any{"field": field}}
+func Exists(field string) Query {
+	return Query{"exists": Query{"field": field}}
 }
 
 // Nested is a shortcut for a nested query
-func Nested(path string, query map[string]any) map[string]any {
-	return map[string]any{"nested": map[string]any{"path": path, "query": query}}
+func Nested(path string, query Query) Query {
+	return Query{"nested": Query{"path": path, "query": query}}
 }
 
 // Match is a shortcut for a match query
-func Match(field string, value any) map[string]any {
-	return map[string]any{"match": map[string]any{field: map[string]any{"query": value}}}
+func Match(field string, value any) Query {
+	return Query{"match": Query{field: Query{"query": value}}}
 }
 
 // MatchPhrase is a shortcut for a match_phrase query
-func MatchPhrase(field, value string) map[string]any {
-	return map[string]any{"match_phrase": map[string]any{field: map[string]any{"query": value}}}
+func MatchPhrase(field, value string) Query {
+	return Query{"match_phrase": Query{field: Query{"query": value}}}
 }
 
 // GreaterThan is a shortcut for a range query where x > value
-func GreaterThan(field string, value any) map[string]any {
-	return map[string]any{"range": map[string]any{field: map[string]any{
+func GreaterThan(field string, value any) Query {
+	return Query{"range": Query{field: Query{
 		"from":          value,
 		"include_lower": false,
 		"include_upper": true,
@@ -56,8 +58,8 @@ func GreaterThan(field string, value any) map[string]any {
 }
 
 // GreaterThanOrEqual is a shortcut for a range query where x >= value
-func GreaterThanOrEqual(field string, value any) map[string]any {
-	return map[string]any{"range": map[string]any{field: map[string]any{
+func GreaterThanOrEqual(field string, value any) Query {
+	return Query{"range": Query{field: Query{
 		"from":          value,
 		"include_lower": true,
 		"include_upper": true,
@@ -66,8 +68,8 @@ func GreaterThanOrEqual(field string, value any) map[string]any {
 }
 
 // LessThan is a shortcut for a range query where x < value
-func LessThan(field string, value any) map[string]any {
-	return map[string]any{"range": map[string]any{field: map[string]any{
+func LessThan(field string, value any) Query {
+	return Query{"range": Query{field: Query{
 		"from":          nil,
 		"include_lower": true,
 		"include_upper": false,
@@ -76,8 +78,8 @@ func LessThan(field string, value any) map[string]any {
 }
 
 // LessThanOrEqual is a shortcut for a range query where x <= value
-func LessThanOrEqual(field string, value any) map[string]any {
-	return map[string]any{"range": map[string]any{field: map[string]any{
+func LessThanOrEqual(field string, value any) Query {
+	return Query{"range": Query{field: Query{
 		"from":          nil,
 		"include_lower": true,
 		"include_upper": true,
@@ -86,8 +88,8 @@ func LessThanOrEqual(field string, value any) map[string]any {
 }
 
 // Between is a shortcut for a range query where from <= x < to
-func Between(field string, from, to any) map[string]any {
-	return map[string]any{"range": map[string]any{field: map[string]any{
+func Between(field string, from, to any) Query {
+	return Query{"range": Query{field: Query{
 		"from":          from,
 		"include_lower": true,
 		"include_upper": false,

diff --git a/utils/elastic/sort.go b/utils/elastic/sort.go
@@ -1,13 +1,15 @@
 package elastic
 
-// Sort is a shortcut for a simple field sort
-func Sort(field string, ascending bool) map[string]any {
-	return map[string]any{field: map[string]any{"order": order(ascending)}}
+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 map[string]any, path string, ascending bool) map[string]any {
-	return map[string]any{field: map[string]any{
+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),
 	}}