maxmind / mmdbwriter

Go library for writing MaxMind DB (mmdb) files
Apache License 2.0
124 stars 28 forks source link

Allow `Tree` to have default insert func #22

Closed UmanShahzad closed 3 years ago

UmanShahzad commented 3 years ago

Right now the Tree.Insert function defaults to the inserter.ReplaceWith strategy. You have to use Tree.InsertFunc to use any other strategy.

It would be nice if we could choose the default for Tree.Insert. Something close to what this diff achieves:

diff --git a/tree.go b/tree.go
index 4e436a8..ce0fca9 100644
--- a/tree.go
+++ b/tree.go
@@ -72,6 +72,11 @@ type Options struct {
        // implementations that do not correctly handle metadata pointers. Its
        // use should primarily be limited to existing database types.
        DisableMetadataPointers bool
+
+       // Inserter is the insert function used when calling `Insert`. It defaults
+       // to `inserter.ReplaceWith`, which replaces any conflicting old value
+       // entirely with the new.
+       Inserter func(value mmdbtype.DataType) (mmdbtype.DataType, error)
 }

 // Tree represents an MaxMind DB search tree.
@@ -88,6 +93,7 @@ type Tree struct {
        treeDepth               int
        // This is set when the tree is finalized
        nodeCount int
+       inserterFunc func(value mmdbtype.DataType) (mmdbtype.DataType, error)
 }

 // New creates a new Tree.
@@ -101,6 +107,7 @@ func New(opts Options) (*Tree, error) {
                ipVersion:               6,
                recordSize:              28,
                root:                    &node{},
+               inserterFunc:            inserter.ReplaceWith,
        }

        if opts.BuildEpoch != 0 {
@@ -123,6 +130,10 @@ func New(opts Options) (*Tree, error) {
                tree.recordSize = opts.RecordSize
        }

+       if opts.Inserter != nil {
+               tree.inserterFunc = opts.Inserter
+       }
+
        switch tree.ipVersion {
        case 6:
                tree.treeDepth = 128
@@ -214,7 +225,7 @@ func Load(path string, opts Options) (*Tree, error) {
 //
 // This is not safe to call from multiple threads.
 func (t *Tree) Insert(network *net.IPNet, value mmdbtype.DataType) error {
-       return t.InsertFunc(network, inserter.ReplaceWith(value))
+       return t.InsertFunc(network, t.inserterFunc(value))
 }

 // InsertFunc will insert the output of the function passed to it. The argument
oschwald commented 3 years ago

If you made a PR, I think it would be accepted. The documentation for Insert would probably need to be updated slightly as well.