Merge pull request #11 from harshit98/sort-clause-support

feat: extend support for sort clauses

Author: gsingh-ds

common.go

@@ -22,17 +22,3 @@ func (source Source) Map() map[string]interface{} {
 	}
 	return m
 }
-
-// Sort represents a list of keys to sort by.
-type Sort []map[string]interface{}
-
-// Order is the ordering for a sort key (ascending, descending).
-type Order string
-
-const (
-	// OrderAsc represents sorting in ascending order.
-	OrderAsc Order = "asc"
-
-	// OrderDesc represents sorting in descending order.
-	OrderDesc Order = "desc"
-)

search.go

@@ -7,8 +7,8 @@ import (
 	"fmt"
 	"time"
 
-	opensearch "github.com/opensearch-project/opensearch-go/v4"
-	opensearchapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
+	"github.com/opensearch-project/opensearch-go/v4"
+	"github.com/opensearch-project/opensearch-go/v4/opensearchapi"
 )
 
 // SearchRequest represents the parameters for an OpenSearch query.
@@ -21,7 +21,7 @@ type SearchRequest struct {
 	postFilter   Mappable
 	query        Mappable
 	size         *uint64
-	sort         Sort
+	sort         []SortOption
 	source       Source
 	timeout      *time.Duration
 	scriptFields []*ScriptField
@@ -63,14 +63,12 @@ func (req *SearchRequest) Size(size uint64) *SearchRequest {
 	return req
 }
 
-// Sort sets how the results should be sorted.
-func (req *SearchRequest) Sort(name string, order Order) *SearchRequest {
-	req.sort = append(req.sort, map[string]interface{}{
-		name: map[string]interface{}{
-			"order": order,
-		},
-	})
-
+// Sort appends one or more sort options.
+// Accepts any type that implements SortOption (field, script, raw)
+func (req *SearchRequest) Sort(opts ...SortOption) *SearchRequest {
+	if opts != nil {
+		req.sort = append(req.sort, opts...)
+	}
 	return req
 }
 
@@ -136,7 +134,11 @@ func (req *SearchRequest) Map() map[string]interface{} {
 		m["size"] = *req.size
 	}
 	if len(req.sort) > 0 {
-		m["sort"] = req.sort
+		sortSlice := make([]any, 0, len(req.sort))
+		for _, params := range req.sort {
+			sortSlice = append(sortSlice, params.Map())
+		}
+		m["sort"] = sortSlice
 	}
 	if req.from != nil {
 		m["from"] = *req.from

search_test.go

@@ -57,8 +57,10 @@ func TestSearchMaps(t *testing.T) {
 				Size(30).
 				From(5).
 				Explain(true).
-				Sort("field_1", OrderDesc).
-				Sort("field_2", OrderAsc).
+				Sort(
+					FieldSort("field_1").Order(OrderDesc),
+					FieldSort("field_2").Order(OrderAsc),
+				).
 				SourceIncludes("field_1", "field_2").
 				SourceExcludes("field_3").
 				Timeout(time.Duration(20000000000)).

sort.go

@@ -0,0 +1,141 @@
+package osquery
+
+// Order is the ordering for a sort key (ascending, descending).
+type Order string
+
+const (
+	// OrderAsc represents sorting in ascending order.
+	OrderAsc Order = "asc"
+
+	// OrderDesc represents sorting in descending order.
+	OrderDesc Order = "desc"
+)
+
+// Mode is the mode for a sort key (min, max, sum, avg, median).
+type Mode string
+
+const (
+	// SortModeMin represents the minimum value.
+	SortModeMin Mode = "min"
+
+	// SortModeMax represents the maximum value.
+	SortModeMax Mode = "max"
+
+	// SortModeSum represents the sum of values.
+	SortModeSum Mode = "sum"
+
+	// SortModeAvg represents the average of values.
+	SortModeAvg Mode = "avg"
+
+	// SortModeMedian represents the median of values.
+	SortModeMedian Mode = "median"
+)
+
+// SortOption is an interface for different types of sort options
+type SortOption interface {
+	Map() map[string]any
+}
+
+// ScriptSortOption represents a script-based sort option for elasticsearch
+type ScriptSortOption struct {
+	sortType string
+	script   *ScriptField
+	order    Order
+}
+
+// ScriptSort creates a new query of type "_script" with the provided
+// type and script.
+func ScriptSort(scriptField *ScriptField, sortType string) *ScriptSortOption {
+	return &ScriptSortOption{
+		script:   scriptField,
+		sortType: sortType,
+	}
+}
+
+func (s *ScriptSortOption) Order(order Order) *ScriptSortOption {
+	s.order = order
+	return s
+}
+
+func (s *ScriptSortOption) Map() map[string]any {
+	scriptMapRaw, ok := s.script.Map()["script"]
+	if !ok {
+		return nil
+	}
+
+	scriptMap, ok := scriptMapRaw.(map[string]any)
+	if !ok {
+		return nil
+	}
+
+	sortOptions := map[string]any{
+		"type":   s.sortType,
+		"script": scriptMap,
+	}
+
+	if s.order != "" {
+		sortOptions["order"] = s.order
+	}
+
+	return map[string]any{
+		"_script": sortOptions,
+	}
+}
+
+type FieldSortOption struct {
+	field        string
+	order        Order
+	mode         Mode
+	nestedPath   string
+	nestedFilter Mappable
+}
+
+func FieldSort(field string) *FieldSortOption {
+	return &FieldSortOption{
+		field: field,
+	}
+}
+
+func (f *FieldSortOption) Order(order Order) *FieldSortOption {
+	f.order = order
+	return f
+}
+
+func (f *FieldSortOption) Mode(mode Mode) *FieldSortOption {
+	f.mode = mode
+	return f
+}
+
+func (f *FieldSortOption) NestedPath(nestedPath string) *FieldSortOption {
+	f.nestedPath = nestedPath
+	return f
+}
+
+func (f *FieldSortOption) NestedFilter(nestedFilter Mappable) *FieldSortOption {
+	f.nestedFilter = nestedFilter
+	return f
+}
+
+func (f *FieldSortOption) Map() map[string]any {
+	sortOptions := map[string]any{}
+
+	if f.order != "" {
+		sortOptions["order"] = f.order
+	}
+
+	if f.mode != "" {
+		sortOptions["mode"] = f.mode
+	}
+
+	if f.nestedPath != "" {
+		sortOptions["nested_path"] = f.nestedPath
+
+		if f.nestedFilter != nil {
+			sortOptions["nested_filter"] = f.nestedFilter.Map()
+		}
+	}
+
+	return map[string]any{
+		f.field: sortOptions,
+	}
+}