diff --git a/aggs_histogram.go b/aggs_histogram.go new file mode 100644 index 0000000..e7c94be --- /dev/null +++ b/aggs_histogram.go @@ -0,0 +1,71 @@ +package osquery + +type HistogramAggregation struct { + name string + field string + interval float64 + offset *float64 + minDocCount *int + aggs []Aggregation +} + +// HistogramAgg creates a new histogram aggregation. +func HistogramAgg(name string, field string, interval float64) *HistogramAggregation { + return &HistogramAggregation{ + name: name, + field: field, + interval: interval, + } +} + +// Name returns the name of the aggregation. +func (agg *HistogramAggregation) Name() string { + return agg.name +} + +// Offset sets an optional offset value. +func (agg *HistogramAggregation) Offset(offset float64) *HistogramAggregation { + agg.offset = &offset + return agg +} + +// MinDocCount sets the optional minimum document count for buckets. +func (agg *HistogramAggregation) MinDocCount(min int) *HistogramAggregation { + agg.minDocCount = &min + return agg +} + +// Aggs sets sub-aggregations for the histogram buckets. +func (agg *HistogramAggregation) Aggs(aggs ...Aggregation) *HistogramAggregation { + agg.aggs = aggs + return agg +} + +// Map builds the OpenSearch aggregation map. +func (agg *HistogramAggregation) Map() map[string]interface{} { + histogramMap := map[string]interface{}{ + "field": agg.field, + "interval": agg.interval, + } + + if agg.offset != nil { + histogramMap["offset"] = *agg.offset + } + if agg.minDocCount != nil { + histogramMap["min_doc_count"] = *agg.minDocCount + } + + outerMap := map[string]interface{}{ + "histogram": histogramMap, + } + + if len(agg.aggs) > 0 { + subAggs := make(map[string]map[string]interface{}) + for _, sub := range agg.aggs { + subAggs[sub.Name()] = sub.Map() + } + outerMap["aggs"] = subAggs + } + + return outerMap +} diff --git a/aggs_histogram_test.go b/aggs_histogram_test.go new file mode 100644 index 0000000..7173335 --- /dev/null +++ b/aggs_histogram_test.go @@ -0,0 +1,50 @@ +package osquery + +import "testing" + +func TestHistogramAggs(t *testing.T) { + runMapTests(t, []mapTest{ + { + "histogram agg: basic", + HistogramAgg("hist", "price", 10), + map[string]interface{}{ + "histogram": map[string]interface{}{ + "field": "price", + "interval": 10.0, + }, + }, + }, + { + "histogram agg: with offset and min_doc_count", + HistogramAgg("hist", "price", 5). + Offset(2). + MinDocCount(1), + map[string]interface{}{ + "histogram": map[string]interface{}{ + "field": "price", + "interval": 5.0, + "offset": 2.0, + "min_doc_count": 1, + }, + }, + }, + { + "histogram agg: with sub-aggs", + HistogramAgg("hist", "price", 20). + Aggs(Avg("avg_price", "price")), + map[string]interface{}{ + "histogram": map[string]interface{}{ + "field": "price", + "interval": 20.0, + }, + "aggs": map[string]interface{}{ + "avg_price": map[string]interface{}{ + "avg": map[string]interface{}{ + "field": "price", + }, + }, + }, + }, + }, + }) +} diff --git a/aggs_reverse_nested.go b/aggs_reverse_nested.go new file mode 100644 index 0000000..5c5a7ce --- /dev/null +++ b/aggs_reverse_nested.go @@ -0,0 +1,53 @@ +package osquery + +type ReverseNestedAggregation struct { + name string + path *string // Optional + aggs []Aggregation // Optional sub-aggregations +} + +// ReverseNestedAgg creates a new reverse_nested aggregation. +func ReverseNestedAgg(name string) *ReverseNestedAggregation { + return &ReverseNestedAggregation{ + name: name, + } +} + +// Name returns the name of the aggregation. +func (agg *ReverseNestedAggregation) Name() string { + return agg.name +} + +// Path sets the optional reverse_nested path. +func (agg *ReverseNestedAggregation) Path(p string) *ReverseNestedAggregation { + agg.path = &p + return agg +} + +// Aggs sets optional sub-aggregations. +func (agg *ReverseNestedAggregation) Aggs(aggs ...Aggregation) *ReverseNestedAggregation { + agg.aggs = aggs + return agg +} + +// Map builds the OpenSearch aggregation map. +func (agg *ReverseNestedAggregation) Map() map[string]interface{} { + reverseNestedBody := make(map[string]interface{}) + if agg.path != nil { + reverseNestedBody["path"] = *agg.path + } + + outerMap := map[string]interface{}{ + "reverse_nested": reverseNestedBody, + } + + if len(agg.aggs) > 0 { + subAggs := make(map[string]map[string]interface{}) + for _, sub := range agg.aggs { + subAggs[sub.Name()] = sub.Map() + } + outerMap["aggs"] = subAggs + } + + return outerMap +} diff --git a/aggs_reverse_nested_test.go b/aggs_reverse_nested_test.go new file mode 100644 index 0000000..0a92aec --- /dev/null +++ b/aggs_reverse_nested_test.go @@ -0,0 +1,39 @@ +package osquery + +import "testing" + +func TestReverseNestedAggs(t *testing.T) { + runMapTests(t, []mapTest{ + { + "reverse_nested agg: basic", + ReverseNestedAgg("to_parent"), + map[string]interface{}{ + "reverse_nested": map[string]interface{}{}, + }, + }, + { + "reverse_nested agg: with path", + ReverseNestedAgg("to_root").Path("some.nested.path"), + map[string]interface{}{ + "reverse_nested": map[string]interface{}{ + "path": "some.nested.path", + }, + }, + }, + { + "reverse_nested agg: with sub-aggregations", + ReverseNestedAgg("to_parent"). + Aggs(Cardinality("product_count", "group_id")), + map[string]interface{}{ + "reverse_nested": map[string]interface{}{}, + "aggs": map[string]interface{}{ + "product_count": map[string]interface{}{ + "cardinality": map[string]interface{}{ + "field": "group_id", + }, + }, + }, + }, + }, + }) +} diff --git a/collapse.go b/collapse.go new file mode 100644 index 0000000..05110e8 --- /dev/null +++ b/collapse.go @@ -0,0 +1,22 @@ +package osquery + +type Collapse struct { + field string + Mappable +} + +func CollapseField(field string) Collapse { + return Collapse{ + field: field, + } +} + +func (c Collapse) Map() map[string]interface{} { + outerMap := make(map[string]interface{}) + if c.field != "" { + outerMap = map[string]interface{}{ + "field": c.field, + } + } + return outerMap +} diff --git a/collapse_test.go b/collapse_test.go new file mode 100644 index 0000000..126df4f --- /dev/null +++ b/collapse_test.go @@ -0,0 +1,15 @@ +package osquery + +import "testing" + +func TestCollapse(t *testing.T) { + runMapTests(t, []mapTest{ + { + "Basic collapse testing", + CollapseField("variant_group.group_id"), + map[string]interface{}{ + "field": "variant_group.group_id", + }, + }, + }) +} diff --git a/common.go b/common.go index aef5705..1a32cac 100644 --- a/common.go +++ b/common.go @@ -12,8 +12,8 @@ type Source struct { } // Map returns a map representation of the Source object. -func (source Source) Map() map[string]interface{} { - m := make(map[string]interface{}) +func (source Source) Map() map[string]any { + m := make(map[string]any) if len(source.includes) > 0 { m["includes"] = source.includes } @@ -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" -) diff --git a/function_score.go b/function_score.go new file mode 100644 index 0000000..1ab6e61 --- /dev/null +++ b/function_score.go @@ -0,0 +1,155 @@ +// Package osquery Modified by harshit98 on 2025-05-10 +// Changes: Added function score support +package osquery + +type FunctionScoreQuery struct { + query Mappable + functions []Function + boostMode string + scoreMode string + maxBoost *float32 + minScore *float32 + boost *float32 +} + +type Function interface { + Map() map[string]interface{} +} + +type RandomScoreFunction struct { + seed *int64 + field string +} + +type ScriptScoreFunction struct { + script *ScriptField +} + +// Additional functions can be added as per usecase in future like: +// - ScriptScoreFunction +// - DecayFunction (with variants for geo, date, numeric) +// - WeightFunction +// +// For ref: https://docs.opensearch.org/docs/latest/query-dsl/compound/function-score/ + +func FunctionScore(query Mappable) *FunctionScoreQuery { + return &FunctionScoreQuery{query: query} +} + +func (q *FunctionScoreQuery) Function(f Function) *FunctionScoreQuery { + q.functions = append(q.functions, f) + return q +} + +func (q *FunctionScoreQuery) BoostMode(boostMode string) *FunctionScoreQuery { + q.boostMode = boostMode + return q +} + +func (q *FunctionScoreQuery) ScoreMode(scoreMode string) *FunctionScoreQuery { + q.scoreMode = scoreMode + return q +} + +func (q *FunctionScoreQuery) MaxBoost(maxBoost float32) *FunctionScoreQuery { + q.maxBoost = &maxBoost + return q +} + +func (q *FunctionScoreQuery) MinScore(minScore float32) *FunctionScoreQuery { + q.minScore = &minScore + return q +} + +func (q *FunctionScoreQuery) Boost(boost float32) *FunctionScoreQuery { + q.boost = &boost + return q +} + +func (q *FunctionScoreQuery) Map() map[string]interface{} { + m := make(map[string]interface{}) + + if q.query != nil { + m["query"] = q.query.Map() + } + + if len(q.functions) > 0 { + funcs := make([]map[string]interface{}, len(q.functions)) + for i, f := range q.functions { + funcs[i] = f.Map() + } + m["functions"] = funcs + } + + if q.boostMode != "" { + m["boost_mode"] = q.boostMode + } + + if q.maxBoost != nil { + m["max_boost"] = *q.maxBoost + } + + if q.scoreMode != "" { + m["score_mode"] = q.scoreMode + } + + if q.minScore != nil { + m["min_score"] = *q.minScore + } + + if q.boost != nil { + m["boost"] = *q.boost + } + + return map[string]interface{}{ + "function_score": m, + } +} + +func RandomScore() *RandomScoreFunction { + return &RandomScoreFunction{} +} + +func (f *RandomScoreFunction) Seed(seed int64) *RandomScoreFunction { + f.seed = &seed + return f +} + +func (f *RandomScoreFunction) Field(field string) *RandomScoreFunction { + f.field = field + return f +} + +func (f *RandomScoreFunction) Map() map[string]interface{} { + m := make(map[string]interface{}) + + if f.seed != nil { + m["seed"] = *f.seed + } + + if f.field != "" { + m["field"] = f.field + } + + return map[string]interface{}{ + "random_score": m, + } +} + +func FunctionScriptScore(script *ScriptField) *ScriptScoreFunction { + return &ScriptScoreFunction{script: script} +} + +func (f *ScriptScoreFunction) Map() map[string]interface{} { + if f.script == nil { + return map[string]interface{}{ + "script_score": map[string]interface{}{}, + } + } + scriptMap := f.script.Map()["script"].(map[string]interface{}) + return map[string]interface{}{ + "script_score": map[string]interface{}{ + "script": scriptMap, + }, + } +} diff --git a/function_score_test.go b/function_score_test.go new file mode 100644 index 0000000..ce68ad9 --- /dev/null +++ b/function_score_test.go @@ -0,0 +1,233 @@ +package osquery + +import ( + "testing" +) + +func TestFunctionScore(t *testing.T) { + runMapTests(t, []mapTest{ + { + "function_score query with random_score function", + FunctionScore(Term("user", "kimchy")). + Function(RandomScore()), + map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "functions": []map[string]interface{}{ + { + "random_score": map[string]interface{}{}, + }, + }, + }, + }, + }, + { + "function_score query with random_score function and boost_mode", + FunctionScore(Term("user", "kimchy")). + Function(RandomScore()). + BoostMode("sum"), + map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "functions": []map[string]interface{}{ + { + "random_score": map[string]interface{}{}, + }, + }, + "boost_mode": "sum", + }, + }, + }, + { + "function_score query with random_score function with seed", + FunctionScore(Term("user", "kimchy")). + Function(RandomScore().Seed(42)), + map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "functions": []map[string]interface{}{ + { + "random_score": map[string]interface{}{ + "seed": int64(42), + }, + }, + }, + }, + }, + }, + { + "function_score query with random_score function with field", + FunctionScore(Term("user", "kimchy")). + Function(RandomScore().Field("_seq_no")), + map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "functions": []map[string]interface{}{ + { + "random_score": map[string]interface{}{ + "field": "_seq_no", + }, + }, + }, + }, + }, + }, + { + "function_score query with multiple functions", + FunctionScore(Term("user", "kimchy")). + Function(RandomScore()). + Function(RandomScore().Seed(123)), + map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "functions": []map[string]interface{}{ + { + "random_score": map[string]interface{}{}, + }, + { + "random_score": map[string]interface{}{ + "seed": int64(123), + }, + }, + }, + }, + }, + }, + { + "function_score query with match_all query", + FunctionScore(MatchAll()). + Function(RandomScore()), + map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "match_all": map[string]interface{}{}, + }, + "functions": []map[string]interface{}{ + { + "random_score": map[string]interface{}{}, + }, + }, + }, + }, + }, + { + "function_score query with script_score function", + FunctionScore(Term("user", "kimchy")). + Function(FunctionScriptScore(Script("my_script").Source("doc['my_field'].value * 2.0"))), + map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "functions": []map[string]interface{}{ + { + "script_score": map[string]interface{}{ + "script": map[string]interface{}{ + "source": "doc['my_field'].value * 2.0", + }, + }, + }, + }, + }, + }, + }, + { + "function_score query with script_score function with params", + FunctionScore(Term("user", "kimchy")). + Function(FunctionScriptScore(Script("my_script"). + Source("doc['my_field'].value * params.factor"). + Params(ScriptParams{"factor": 2.0}). + Lang("painless"))), + map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "functions": []map[string]interface{}{ + { + "script_score": map[string]interface{}{ + "script": map[string]interface{}{ + "source": "doc['my_field'].value * params.factor", + "params": map[string]interface{}{ + "factor": 2.0, + }, + "lang": "painless", + }, + }, + }, + }, + }, + }, + }, + }) +} + +func TestFunctionScoreWithQuery(t *testing.T) { + runMapTests(t, []mapTest{ + { + "query with function_score", + Query( + FunctionScore(Term("user", "kimchy")). + Function(RandomScore()). + BoostMode("sum"), + ), + map[string]interface{}{ + "query": map[string]interface{}{ + "function_score": map[string]interface{}{ + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "functions": []map[string]interface{}{ + { + "random_score": map[string]interface{}{}, + }, + }, + "boost_mode": "sum", + }, + }, + }, + }, + }) +} diff --git a/query_nested.go b/query_nested.go new file mode 100644 index 0000000..bfd44fe --- /dev/null +++ b/query_nested.go @@ -0,0 +1,72 @@ +// Package osquery Modified by bforbhartiii on 2025-02-24 +// Changes: Added nested query support +package osquery + +import "github.com/fatih/structs" + +type ScoreModeType string + +const ( + // ScoreModeAvg Uses the average relevance score of all matching inner documents + ScoreModeAvg ScoreModeType = "avg" + + // ScoreModeMax Assigns the highest relevance score from the matching inner documents to the parent. + ScoreModeMax ScoreModeType = "max" + + // ScoreModeMin Assigns the lowest relevance score from the matching inner documents to the parent. + ScoreModeMin ScoreModeType = "min" + + // ScoreModeSum Sums the relevance scores of all matching inner documents. + ScoreModeSum ScoreModeType = "sum" + + // ScoreModeNone Ignores the relevance scores of inner documents and assigns a score of 0 to the parent document. + ScoreModeNone ScoreModeType = "none" +) + +// NestedQuery represents a compound query of type "nested", +// as described in https://opensearch.org/docs/latest/query-dsl/joining/nested/ +type NestedQuery struct { + path string + query Mappable + name string + scoreMode string + innerHits map[string]interface{} +} + +// Nested creates a new query of type "nested" with the provided path and query. +func Nested(path string, query Mappable) *NestedQuery { + return &NestedQuery{ + path: path, + query: query, + } +} + +// ScoreMode sets the score mode of the query. +func (q *NestedQuery) ScoreMode(mode ScoreModeType) *NestedQuery { + q.scoreMode = string(mode) + return q +} + +// InnerHits sets the inner_hits field of the query. +func (q *NestedQuery) InnerHits(innerHits map[string]interface{}) *NestedQuery { + q.innerHits = innerHits + return q +} + +func (q *NestedQuery) Name(name string) *NestedQuery { + q.name = name + return q +} + +// Map returns a map representation of the query, implementing the Mappable interface. +func (q *NestedQuery) Map() map[string]interface{} { + return map[string]interface{}{ + "nested": structs.Map(struct { + Path string `structs:"path"` + Query map[string]interface{} `structs:"query"` + Name string `structs:"_name,omitempty"` + ScoreMode string `structs:"score_mode,omitempty"` + InnerHits map[string]interface{} `structs:"inner_hits,omitempty"` + }{q.path, q.query.Map(), q.name, q.scoreMode, q.innerHits}), + } +} diff --git a/query_nested_test.go b/query_nested_test.go new file mode 100644 index 0000000..952de5e --- /dev/null +++ b/query_nested_test.go @@ -0,0 +1,78 @@ +package osquery + +import "testing" + +// Test cases for NestedQuery +func TestNestedQuery(t *testing.T) { + runMapTests(t, []mapTest{ + { + "nested query without optional fields", + Nested("comments", Term("user", "kimchy")), + map[string]interface{}{ + "nested": map[string]interface{}{ + "path": "comments", + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + }, + }, + }, + { + "nested query with score_mode", + Nested("comments", Term("user", "kimchy")).ScoreMode(ScoreModeMax), + map[string]interface{}{ + "nested": map[string]interface{}{ + "path": "comments", + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "score_mode": "max", + }, + }, + }, + { + "nested query with inner_hits", + Nested("comments", Term("user", "kimchy")).InnerHits(map[string]interface{}{"size": 3}), + map[string]interface{}{ + "nested": map[string]interface{}{ + "path": "comments", + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "inner_hits": map[string]interface{}{ + "size": 3, + }, + }, + }, + }, + { + "nested query with name", + Nested("comments", Term("user", "kimchy")).Name("nested_comments"), + map[string]interface{}{ + "nested": map[string]interface{}{ + "path": "comments", + "query": map[string]interface{}{ + "term": map[string]interface{}{ + "user": map[string]interface{}{ + "value": "kimchy", + }, + }, + }, + "name": "nested_comments", + }, + }, + }, + }) +} diff --git a/query_script_score.go b/query_script_score.go index b74766f..2d401fd 100644 --- a/query_script_score.go +++ b/query_script_score.go @@ -1,6 +1,5 @@ // Modified by Sushmita on 2025-01-31 -// Changes: Updated ElasticSearch client to OpenSearch client, changed package name to 'osquery', -// updated references to OpenSearch documentation, and modified examples accordingly. +// Changes: Added script score query support package osquery diff --git a/query_term_level.go b/query_term_level.go index 2c995fb..5eef735 100644 --- a/query_term_level.go +++ b/query_term_level.go @@ -448,6 +448,7 @@ type TermsQuery struct { field string values []interface{} boost float32 + name string } // Terms creates a new query of type "terms" on the provided field, and @@ -465,6 +466,12 @@ func (q *TermsQuery) Values(values ...interface{}) *TermsQuery { return q } +// Name sets the name for the query. +func (q *TermsQuery) Name(name string) *TermsQuery { + q.name = name + return q +} + // Boost sets the boost value of the query. func (q *TermsQuery) Boost(b float32) *TermsQuery { q.boost = b @@ -478,6 +485,9 @@ func (q TermsQuery) Map() map[string]interface{} { if q.boost > 0 { innerMap["boost"] = q.boost } + if q.name != "" { + innerMap["_name"] = q.name + } return map[string]interface{}{"terms": innerMap} } diff --git a/query_term_level_test.go b/query_term_level_test.go index 461e857..5b3c92c 100644 --- a/query_term_level_test.go +++ b/query_term_level_test.go @@ -140,6 +140,17 @@ func TestTermLevel(t *testing.T) { }, }, }, + { + "terms", + Terms("user").Values("bla", "pl").Boost(1.3).Name("test"), + map[string]interface{}{ + "terms": map[string]interface{}{ + "user": []string{"bla", "pl"}, + "boost": 1.3, + "_name": "test", + }, + }, + }, { "terms_set", TermsSet("programming_languages", "go", "rust", "COBOL").MinimumShouldMatchField("required_matches"), diff --git a/search.go b/search.go index 0150469..dd40ec9 100644 --- a/search.go +++ b/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. @@ -20,8 +20,9 @@ type SearchRequest struct { searchAfter []interface{} postFilter Mappable query Mappable + collapse Collapse size *uint64 - sort Sort + sort []SortOption source Source timeout *time.Duration scriptFields []*ScriptField @@ -63,14 +64,18 @@ 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, - }, - }) +// Collapse sets field collapsing for the request. +func (req *SearchRequest) Collapse(collapse Collapse) *SearchRequest { + req.collapse = collapse + return req +} +// 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 } @@ -135,8 +140,18 @@ func (req *SearchRequest) Map() map[string]interface{} { if req.size != nil { m["size"] = *req.size } + + collapse := req.collapse.Map() + if len(collapse) > 0 { + m["collapse"] = collapse + } + 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 diff --git a/search_test.go b/search_test.go index 403b988..2a9529c 100644 --- a/search_test.go +++ b/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)). @@ -149,5 +151,14 @@ func TestSearchMaps(t *testing.T) { }, }, }, + { + "a search with collapse", + Search().Collapse(CollapseField("variant_group.group_id")), + map[string]interface{}{ + "collapse": map[string]interface{}{ + "field": "variant_group.group_id", + }, + }, + }, }) } diff --git a/sort.go b/sort.go new file mode 100644 index 0000000..7809d9c --- /dev/null +++ b/sort.go @@ -0,0 +1,160 @@ +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 + GetOrder() Order +} + +// 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) GetOrder() Order { + return s.order +} + +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 + missing string + 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) GetOrder() Order { + return f.order +} + +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) Missing(missing string) *FieldSortOption { + f.missing = missing + 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.missing != "" { + sortOptions["missing"] = f.missing + } + + 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, + } +} diff --git a/sort_test.go b/sort_test.go new file mode 100644 index 0000000..821a432 --- /dev/null +++ b/sort_test.go @@ -0,0 +1,257 @@ +// Package osquery Modified by harshit98 on 2025-05-07 +// Changes: Added sort params support like mode, nested_path, nested_filter +package osquery + +import ( + "testing" +) + +func TestSortExtensions(t *testing.T) { + fieldSortWithOrder := FieldSort("field").Order(OrderAsc) + fieldSortWithOrderAndMode := FieldSort("field").Order(OrderDesc).Mode(SortModeAvg) + + nestedFieldSort := FieldSort("nested.field").Order(OrderAsc).NestedPath("nested") + + nestedFieldSortWithFilter := FieldSort("nested.field"). + Order(OrderAsc). + NestedPath("nested"). + NestedFilter(Match("nested.type").Query("value")) + + nestedFieldSortWithOrderAndMode := FieldSort("nested.field"). + Order(OrderDesc). + Mode(SortModeMax). + NestedPath("nested"). + NestedFilter(Match("nested.type").Query("value")) + + multipleSort1 := FieldSort("field1").Order(OrderAsc) + + multipleSort2 := FieldSort("nested.field"). + Order(OrderDesc). + Mode(SortModeMin). + NestedPath("nested"). + NestedFilter(Match("nested.type").Query("value")) + + runMapTests(t, []mapTest{ + { + "sort with basic order only", + Search().Sort(fieldSortWithOrder), + map[string]any{ + "sort": []map[string]any{ + { + "field": map[string]any{ + "order": "asc", + }, + }, + }, + }, + }, + { + "sort with mode", + Search().Sort(fieldSortWithOrderAndMode), + map[string]any{ + "sort": []map[string]any{ + { + "field": map[string]any{ + "order": "desc", + "mode": "avg", + }, + }, + }, + }, + }, + { + "sort with nested_path", + Search().Sort(nestedFieldSort), + map[string]any{ + "sort": []map[string]any{ + { + "nested.field": map[string]any{ + "order": "asc", + "nested_path": "nested", + }, + }, + }, + }, + }, + { + "sort with nested_path and nested_filter", + Search().Sort(nestedFieldSortWithFilter), + map[string]any{ + "sort": []map[string]any{ + { + "nested.field": map[string]any{ + "order": "asc", + "nested_path": "nested", + "nested_filter": map[string]any{ + "match": map[string]any{ + "nested.type": map[string]any{ + "query": "value", + }, + }, + }, + }, + }, + }, + }, + }, + { + "sort with mode, nested_path and nested_filter", + Search().Sort(nestedFieldSortWithOrderAndMode), + map[string]any{ + "sort": []map[string]any{ + { + "nested.field": map[string]any{ + "order": "desc", + "mode": "max", + "nested_path": "nested", + "nested_filter": map[string]any{ + "match": map[string]any{ + "nested.type": map[string]any{ + "query": "value", + }, + }, + }, + }, + }, + }, + }, + }, + { + "multiple sorts with different options", + Search().Sort(multipleSort1, multipleSort2), + map[string]any{ + "sort": []map[string]any{ + { + "field1": map[string]any{ + "order": "asc", + }, + }, + { + "nested.field": map[string]any{ + "order": "desc", + "mode": "min", + "nested_path": "nested", + "nested_filter": map[string]any{ + "match": map[string]any{ + "nested.type": map[string]any{ + "query": "value", + }, + }, + }, + }, + }, + }, + }, + }, + }) +} + +func TestScriptSortExtensions(t *testing.T) { + // Create script fields for reuse + scriptFieldSort1 := Script("test_script"). + Source("doc['field_name'].value"). + Lang("painless") + + scriptFieldSort2 := Script("test_script"). + Source("doc['field_name'].value * params.factor"). + Lang("painless"). + Params(ScriptParams{"factor": 1.5}) + + scriptFieldSort3 := Script("test_script"). + Source("if (doc['parent_obj.score_field'].size()!=0) { return ( Math.log(doc['parent_obj.score_field'].value*100 + 10 ) * _score ) } else { return _score }"). + Lang("painless") + + // Create script sort params for reuse + scriptSortParams1 := ScriptSort(scriptFieldSort1, "number").Order(OrderDesc) + scriptSortParams2 := ScriptSort(scriptFieldSort2, "number").Order(OrderAsc) + scriptSortParams3 := ScriptSort(scriptFieldSort3, "number").Order(OrderDesc) + + docScoreFieldSort := FieldSort("_score") + regularFieldSort := FieldSort("regular_field").Order(OrderAsc) + + runMapTests(t, []mapTest{ + { + "sort with script", + Search().Sort(scriptSortParams1), + map[string]any{ + "sort": []map[string]any{ + { + "_script": map[string]any{ + "type": "number", + "script": map[string]any{ + "source": "doc['field_name'].value", + "lang": "painless", + }, + "order": "desc", + }, + }, + }, + }, + }, + { + "sort with script and params", + Search().Sort(scriptSortParams2), + map[string]any{ + "sort": []map[string]any{ + { + "_script": map[string]any{ + "type": "number", + "script": map[string]any{ + "source": "doc['field_name'].value * params.factor", + "lang": "painless", + "params": map[string]any{ + "factor": 1.5, + }, + }, + "order": "asc", + }, + }, + }, + }, + }, + { + "sort with raw field and script", + Search().Sort(docScoreFieldSort, scriptSortParams3), + map[string]any{ + "sort": []any{ + map[string]any{ + "_score": map[string]any{}, + }, + map[string]any{ + "_script": map[string]any{ + "type": "number", + "script": map[string]any{ + "source": "if (doc['parent_obj.score_field'].size()!=0) { return ( Math.log(doc['parent_obj.score_field'].value*100 + 10 ) * _score ) } else { return _score }", + "lang": "painless", + }, + "order": "desc", + }, + }, + }, + }, + }, + { + "mixed sort with field and script", + Search().Sort(regularFieldSort, scriptSortParams1), + map[string]any{ + "sort": []map[string]any{ + { + "regular_field": map[string]any{ + "order": "asc", + }, + }, + { + "_script": map[string]any{ + "type": "number", + "script": map[string]any{ + "source": "doc['field_name'].value", + "lang": "painless", + }, + "order": "desc", + }, + }, + }, + }, + }, + }) +}