Skip to content

Commit

Permalink
Add IPRange aggregation
Browse files Browse the repository at this point in the history
This commit adds the IP range aggregation type described in [1].

To access the buckets in response, use `IPRange`.

[1] https://www.elastic.co/guide/en/elasticsearch/reference/6.0/search-aggregations-bucket-iprange-aggregation.html
  • Loading branch information
olivere committed Nov 16, 2017
1 parent 30c7022 commit 17485ea
Show file tree
Hide file tree
Showing 4 changed files with 289 additions and 4 deletions.
4 changes: 2 additions & 2 deletions search_aggs.go
Original file line number Diff line number Diff line change
Expand Up @@ -398,9 +398,9 @@ func (a Aggregations) DateRange(name string) (*AggregationBucketRangeItems, bool
return nil, false
}

// IPv4Range returns IPv4 range aggregation results.
// IPRange returns IP range aggregation results.
// See: https://www.elastic.co/guide/en/elasticsearch/reference/6.0/search-aggregations-bucket-iprange-aggregation.html
func (a Aggregations) IPv4Range(name string) (*AggregationBucketRangeItems, bool) {
func (a Aggregations) IPRange(name string) (*AggregationBucketRangeItems, bool) {
if raw, found := a[name]; found {
agg := new(AggregationBucketRangeItems)
if raw == nil {
Expand Down
195 changes: 195 additions & 0 deletions search_aggs_bucket_ip_range.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,195 @@
// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

// IPRangeAggregation is a range aggregation that is dedicated for
// IP addresses.
//
// See: https://www.elastic.co/guide/en/elasticsearch/reference/6.0/search-aggregations-bucket-iprange-aggregation.html
type IPRangeAggregation struct {
field string
subAggregations map[string]Aggregation
meta map[string]interface{}
keyed *bool
entries []IPRangeAggregationEntry
}

type IPRangeAggregationEntry struct {
Key string
Mask string
From string
To string
}

func NewIPRangeAggregation() *IPRangeAggregation {
return &IPRangeAggregation{
subAggregations: make(map[string]Aggregation),
entries: make([]IPRangeAggregationEntry, 0),
}
}

func (a *IPRangeAggregation) Field(field string) *IPRangeAggregation {
a.field = field
return a
}

func (a *IPRangeAggregation) SubAggregation(name string, subAggregation Aggregation) *IPRangeAggregation {
a.subAggregations[name] = subAggregation
return a
}

// Meta sets the meta data to be included in the aggregation response.
func (a *IPRangeAggregation) Meta(metaData map[string]interface{}) *IPRangeAggregation {
a.meta = metaData
return a
}

func (a *IPRangeAggregation) Keyed(keyed bool) *IPRangeAggregation {
a.keyed = &keyed
return a
}

func (a *IPRangeAggregation) AddMaskRange(mask string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{Mask: mask})
return a
}

func (a *IPRangeAggregation) AddMaskRangeWithKey(key, mask string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, Mask: mask})
return a
}

func (a *IPRangeAggregation) AddRange(from, to string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{From: from, To: to})
return a
}

func (a *IPRangeAggregation) AddRangeWithKey(key, from, to string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: from, To: to})
return a
}

func (a *IPRangeAggregation) AddUnboundedTo(from string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{From: from, To: ""})
return a
}

func (a *IPRangeAggregation) AddUnboundedToWithKey(key, from string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: from, To: ""})
return a
}

func (a *IPRangeAggregation) AddUnboundedFrom(to string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{From: "", To: to})
return a
}

func (a *IPRangeAggregation) AddUnboundedFromWithKey(key, to string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: "", To: to})
return a
}

func (a *IPRangeAggregation) Lt(to string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{From: "", To: to})
return a
}

func (a *IPRangeAggregation) LtWithKey(key, to string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: "", To: to})
return a
}

func (a *IPRangeAggregation) Between(from, to string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{From: from, To: to})
return a
}

func (a *IPRangeAggregation) BetweenWithKey(key, from, to string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: from, To: to})
return a
}

func (a *IPRangeAggregation) Gt(from string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{From: from, To: ""})
return a
}

func (a *IPRangeAggregation) GtWithKey(key, from string) *IPRangeAggregation {
a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: from, To: ""})
return a
}

func (a *IPRangeAggregation) Source() (interface{}, error) {
// Example:
// {
// "aggs" : {
// "range" : {
// "ip_range": {
// "field": "ip",
// "ranges": [
// { "to": "10.0.0.5" },
// { "from": "10.0.0.5" }
// ]
// }
// }
// }
// }
// }
//
// This method returns only the { "ip_range" : { ... } } part.

source := make(map[string]interface{})
opts := make(map[string]interface{})
source["ip_range"] = opts

// ValuesSourceAggregationBuilder
if a.field != "" {
opts["field"] = a.field
}

if a.keyed != nil {
opts["keyed"] = *a.keyed
}

var ranges []interface{}
for _, ent := range a.entries {
r := make(map[string]interface{})
if ent.Key != "" {
r["key"] = ent.Key
}
if ent.Mask != "" {
r["mask"] = ent.Mask
} else {
if ent.From != "" {
r["from"] = ent.From
}
if ent.To != "" {
r["to"] = ent.To
}
}
ranges = append(ranges, r)
}
opts["ranges"] = ranges

// AggregationBuilder (SubAggregations)
if len(a.subAggregations) > 0 {
aggsMap := make(map[string]interface{})
source["aggregations"] = aggsMap
for name, aggregate := range a.subAggregations {
src, err := aggregate.Source()
if err != nil {
return nil, err
}
aggsMap[name] = src
}
}

// Add Meta data if available
if len(a.meta) > 0 {
source["meta"] = a.meta
}

return source, nil
}
90 changes: 90 additions & 0 deletions search_aggs_bucket_ip_range_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

import (
"encoding/json"
"testing"
)

func TestIPRangeAggregation(t *testing.T) {
agg := NewIPRangeAggregation().Field("remote_ip")
agg = agg.AddRange("", "10.0.0.0")
agg = agg.AddRange("10.1.0.0", "10.1.255.255")
agg = agg.AddRange("10.2.0.0", "")
src, err := agg.Source()
if err != nil {
t.Fatal(err)
}
data, err := json.Marshal(src)
if err != nil {
t.Fatalf("marshaling to JSON failed: %v", err)
}
got := string(data)
expected := `{"ip_range":{"field":"remote_ip","ranges":[{"to":"10.0.0.0"},{"from":"10.1.0.0","to":"10.1.255.255"},{"from":"10.2.0.0"}]}}`
if got != expected {
t.Errorf("expected\n%s\n,got:\n%s", expected, got)
}
}

func TestIPRangeAggregationMask(t *testing.T) {
agg := NewIPRangeAggregation().Field("remote_ip")
agg = agg.AddMaskRange("10.0.0.0/25")
agg = agg.AddMaskRange("10.0.0.127/25")
src, err := agg.Source()
if err != nil {
t.Fatal(err)
}
data, err := json.Marshal(src)
if err != nil {
t.Fatalf("marshaling to JSON failed: %v", err)
}
got := string(data)
expected := `{"ip_range":{"field":"remote_ip","ranges":[{"mask":"10.0.0.0/25"},{"mask":"10.0.0.127/25"}]}}`
if got != expected {
t.Errorf("expected\n%s\n,got:\n%s", expected, got)
}
}

func TestIPRangeAggregationWithKeyedFlag(t *testing.T) {
agg := NewIPRangeAggregation().Field("remote_ip")
agg = agg.Keyed(true)
agg = agg.AddRange("", "10.0.0.0")
agg = agg.AddRange("10.1.0.0", "10.1.255.255")
agg = agg.AddRange("10.2.0.0", "")
src, err := agg.Source()
if err != nil {
t.Fatal(err)
}
data, err := json.Marshal(src)
if err != nil {
t.Fatalf("marshaling to JSON failed: %v", err)
}
got := string(data)
expected := `{"ip_range":{"field":"remote_ip","keyed":true,"ranges":[{"to":"10.0.0.0"},{"from":"10.1.0.0","to":"10.1.255.255"},{"from":"10.2.0.0"}]}}`
if got != expected {
t.Errorf("expected\n%s\n,got:\n%s", expected, got)
}
}

func TestIPRangeAggregationWithKeys(t *testing.T) {
agg := NewIPRangeAggregation().Field("remote_ip")
agg = agg.Keyed(true)
agg = agg.LtWithKey("infinity", "10.0.0.5")
agg = agg.GtWithKey("and-beyond", "10.0.0.5")
src, err := agg.Source()
if err != nil {
t.Fatal(err)
}
data, err := json.Marshal(src)
if err != nil {
t.Fatalf("marshaling to JSON failed: %v", err)
}
got := string(data)
expected := `{"ip_range":{"field":"remote_ip","keyed":true,"ranges":[{"key":"infinity","to":"10.0.0.5"},{"from":"10.0.0.5","key":"and-beyond"}]}}`
if got != expected {
t.Errorf("expected\n%s\n,got:\n%s", expected, got)
}
}
4 changes: 2 additions & 2 deletions search_aggs_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2420,7 +2420,7 @@ func TestAggsBucketDateRange(t *testing.T) {
}
}

func TestAggsBucketIPv4Range(t *testing.T) {
func TestAggsBucketIPRange(t *testing.T) {
s := `{
"ip_ranges": {
"buckets" : [
Expand All @@ -2444,7 +2444,7 @@ func TestAggsBucketIPv4Range(t *testing.T) {
t.Fatalf("expected no error decoding; got: %v", err)
}

agg, found := aggs.IPv4Range("ip_ranges")
agg, found := aggs.IPRange("ip_ranges")
if !found {
t.Fatalf("expected aggregation to be found; got: %v", found)
}
Expand Down

0 comments on commit 17485ea

Please sign in to comment.